Fix LSP latency, mutex contention, and session correctness

baml_language/crates/baml_builtins2/baml_std/testing/registry.baml

@@ -128,6 +128,35 @@ class TestRegistry {
     return items
   }
 
+  function is_expanded(self, name: string) -> bool {
+    if (self.expansions.has(name)) {
+      return true
+    }
+    for (let key in self.expansions.keys()) {
+      if (name.starts_with(key + "/")) {
+        return self.expansions[key].is_expanded(name)
+      }
+    }
+    return false
+  }
+
+  // Roll back one newly-created expansion after serialization/publication
+  // failure. The Rust caller serializes expansion mutation with the registry
+  // owner and invokes this only when `is_expanded(name)` was initially false.
+  function rollback_expand_set(self, name: string) -> null {
+    if (self.expansions.has(name)) {
+      self.expansions.delete(name)
+      return null
+    }
+    for (let key in self.expansions.keys()) {
+      if (name.starts_with(key + "/")) {
+        self.expansions[key].rollback_expand_set(name)
+        return null
+      }
+    }
+    return null
+  }
+
   // Find test by name in root collector, then recurse into expanded sub-registries.
   function run_test(self, name: string) -> TestReport {
     // Check tests in this registry's collector
@@ -192,18 +221,25 @@ class TestRegistry {
   // Expand a lazy testset by name. Searches recursively through expanded
   // sub-registries so nested testsets can be expanded too.
   function expand_set(self, name: string) -> SerializedTestDef[] {
+    self.expand_set_in_place(name)
+    return self.serialize()
+  }
+
+  // Mutation-only entry point used by the playground's transactional
+  // expand/serialize/publish path.
+  function expand_set_in_place(self, name: string) -> null {
     // Try root collector first
     for (let ts in self.collector.testsets) {
       if (ts.name == name) {
         self.expand_testset_registry(ts)
-        return self.serialize()
+        return null
       }
     }
     // Search in expanded sub-registries using segment-aware prefix match
     for (let k in self.expansions.keys()) {
       let sub = self.expansions[k]
       if (name.starts_with(k + "/")) {
-        return sub.expand_set(name)
+        return sub.expand_set_in_place(name)
       }
     }
     throw ("TestSet not found for expansion: " + name)

baml_language/crates/baml_cli/src/paint.rs

@@ -8,11 +8,10 @@
 //! so classification is effectively free (it reads Salsa-cached queries) and can
 //! never drift from the language the way a hand-maintained grammar would.
 //!
-//! [`Highlighter`] caches the per-file token set so a single description (body +
-//! dependency/reference rows, which can touch several files) computes each file's
-//! tokens at most once.
+//! [`Highlighter`] reuses the Salsa-cached per-file token set across a single
+//! description (body + dependency/reference rows, which can touch several files).
 
-use std::{cell::RefCell, collections::HashMap, fmt::Write, path::Path, rc::Rc};
+use std::{fmt::Write, path::Path};
 
 use baml_db::{
     FileId, SourceFile,
@@ -467,31 +466,19 @@ impl Painter {
 
 // ── Source highlighting ─────────────────────────────────────────────────────────
 
-/// Highlights source ranges using the compiler's semantic tokens, caching the
-/// token set per file so repeated lookups within one render are cheap.
+/// Highlights source ranges using the compiler's Salsa-cached semantic tokens.
 pub struct Highlighter<'db> {
     db: &'db ProjectDatabase,
-    cache: RefCell<HashMap<SourceFile, Rc<[SemanticToken]>>>,
 }
 
 impl<'db> Highlighter<'db> {
     pub fn new(db: &'db ProjectDatabase) -> Self {
-        Self {
-            db,
-            cache: RefCell::new(HashMap::new()),
-        }
+        Self { db }
     }
 
-    /// Semantic tokens for `file`, computed once and cached (sorted by start).
-    fn tokens(&self, file: SourceFile) -> Rc<[SemanticToken]> {
-        if let Some(cached) = self.cache.borrow().get(&file) {
-            return Rc::clone(cached);
-        }
-        let mut toks = semantic_tokens(self.db, file);
-        toks.sort_by_key(|t| t.range.start());
-        let rc: Rc<[SemanticToken]> = toks.into();
-        self.cache.borrow_mut().insert(file, Rc::clone(&rc));
-        rc
+    /// Semantic tokens for `file`, memoized by Salsa and emitted in source order.
+    fn tokens(&self, file: SourceFile) -> &[SemanticToken] {
+        semantic_tokens(self.db, file)
     }
 
     /// Highlight the verbatim source slice `file.text()[range]`.
@@ -509,7 +496,7 @@ impl<'db> Highlighter<'db> {
 
         let mut out = String::new();
         let mut cursor = 0usize;
-        for tok in self.tokens(file).iter() {
+        for tok in self.tokens(file) {
             let ts: usize = tok.range.start().into();
             let te: usize = tok.range.end().into();
             if te <= start {

baml_language/crates/baml_lsp2_actions/src/annotations.rs

@@ -1,7 +1,7 @@
 //! Inline type / parameter-name annotations for BAML files (inlay hints).
 //!
-//! Provides `annotations(db, file) -> Vec<InlineAnnotation>` — a regular
-//! function (not a Salsa query) that walks expression-body functions in a file
+//! Provides `annotations(db, file) -> &Vec<InlineAnnotation>` — a Salsa-tracked
+//! query that walks expression-body functions in a file
 //! (top-level functions, class/interface methods, and the synthesized
 //! `$init_test` registration functions), recursing into lambda bodies (e.g.
 //! the bodies of `test` / `testset` blocks, which lower to lambdas passed to
@@ -81,7 +81,7 @@ pub enum AnnotationKind {
 }
 
 /// A single inline annotation (inlay hint) to display in the editor.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)]
 pub struct InlineAnnotation {
     /// Byte offset in the file where the hint is inserted.
     pub offset: TextSize,
@@ -102,9 +102,10 @@ pub struct InlineAnnotation {
 /// Returns annotations sorted in document order (required by the LSP
 /// `textDocument/inlayHint` contract).
 ///
-/// Regular function (not a Salsa query). Internally calls Salsa-cached
-/// queries (`function_body`, `function_body_source_map`,
-/// `infer_scope_types`, `file_item_tree`, `file_semantic_index`).
+/// Salsa-tracked per-file query. Internally calls other Salsa-cached queries
+/// (`function_body`, `function_body_source_map`, `infer_scope_types`,
+/// `file_item_tree`, `file_semantic_index`).
+#[salsa::tracked(returns(ref))]
 pub fn annotations(db: &dyn Db, file: SourceFile) -> Vec<InlineAnnotation> {
     let item_tree = baml_compiler2_hir::file_item_tree(db, file);
     let index = baml_compiler2_hir::file_semantic_index(db, file);

baml_language/crates/baml_lsp2_actions/src/lib.rs

@@ -1,8 +1,9 @@
 //! `baml_lsp2_actions` — IDE action layer built on top of compiler2.
 //!
-//! Modeled after ruff's `ty_ide` crate: regular functions (not Salsa queries)
-//! that take `&dyn Db` and return domain types. Internally they call Salsa
-//! queries from `baml_compiler2_hir` and `baml_compiler2_tir` for cached data.
+//! Modeled after ruff's `ty_ide` crate: IDE functions take `&dyn Db` and return
+//! domain types. Most are regular functions that call cached compiler queries;
+//! frequently repeated whole-file products such as outlines, annotations, and
+//! semantic tokens are Salsa-tracked directly.
 //!
 //! ## Phase 1
 //!
@@ -88,7 +89,7 @@ pub trait Db: baml_compiler2_tir::Db {}
 // ── Public API re-exports ─────────────────────────────────────────────────────
 
 pub use actions::{FileAction, FileActionKind, file_actions};
-pub use annotations::{AnnotationKind, InlineAnnotation, annotations};
+pub use annotations::{AnnotationKind, InlineAnnotation};
 // Re-export `DefinitionKind` so callers (e.g. bex_project) don't need to
 // depend on `baml_compiler2_hir` directly just for type conversions.
 pub use baml_compiler2_hir::contributions::DefinitionKind;
@@ -110,3 +111,12 @@ pub use search::{SymbolInfo, search_symbols};
 pub use tokens::{SemanticToken, SemanticTokenType, TOKEN_TYPES, semantic_tokens};
 pub use type_info::{FunctionParamInfo, TypeInfo, type_at};
 pub use usages::usages_at;
+
+/// Return the Salsa-memoized inline annotations for `file`.
+///
+/// This wrapper preserves the crate-root API without re-exporting Salsa's
+/// generated query type into the same namespace as the public `annotations`
+/// module.
+pub fn annotations(db: &dyn Db, file: baml_base::SourceFile) -> &Vec<InlineAnnotation> {
+    annotations::annotations(db, file)
+}

baml_language/crates/baml_lsp2_actions/src/tokens.rs

@@ -1,7 +1,7 @@
 //! Semantic tokens for BAML files (compiler2 / `lsp2_actions` version).
 //!
-//! Provides `semantic_tokens(db, file) -> Vec<SemanticToken>` using a hybrid
-//! CST + compiler2 approach:
+//! Provides `semantic_tokens(db, file) -> &Vec<SemanticToken>` as a
+//! Salsa-tracked query using a hybrid CST + compiler2 approach:
 //!
 //! - **Structural tokens** (keywords, comments, strings, numbers, operators)
 //!   come from a single CST walk with syntactic classification.
@@ -139,7 +139,7 @@ impl SemanticTokenType {
 // ── SemanticToken ─────────────────────────────────────────────────────────────
 
 /// A classified token ready for LSP encoding.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)]
 pub struct SemanticToken {
     pub range: TextRange,
     pub token_type: SemanticTokenType,
@@ -152,9 +152,10 @@ pub struct SemanticToken {
 /// Always returns tokens in document order (required by the LSP
 /// `textDocument/semanticTokens/full` contract).
 ///
-/// Regular function (not a Salsa query). Internally calls Salsa-cached queries
+/// Salsa-tracked per-file query. Internally calls other Salsa-cached queries
 /// (`infer_scope_types`, `function_body`, `function_body_source_map`,
 /// `file_semantic_index`, `syntax_tree`).
+#[salsa::tracked(returns(ref))]
 pub fn semantic_tokens(db: &dyn Db, file: SourceFile) -> Vec<SemanticToken> {
     let root = baml_compiler_parser::syntax_tree(db, file);
     let mut out = Vec::new();

baml_language/crates/baml_lsp2_actions_tests/Cargo.toml

@@ -16,6 +16,9 @@ baml_project = { workspace = true }
 indexmap = { workspace = true }
 text-size = { workspace = true }
 
+[dev-dependencies]
+salsa = { workspace = true }
+
 [build-dependencies]
 prettyplease = { workspace = true }
 proc-macro2 = { workspace = true }

baml_language/crates/baml_lsp2_actions_tests/src/lib.rs

@@ -2,6 +2,9 @@ pub mod parser;
 pub mod runner;
 pub mod updater;
 
+#[cfg(test)]
+mod memoization;
+
 #[cfg(test)]
 mod test_files {
     include!(concat!(env!("OUT_DIR"), "/generated_lsp2_tests.rs"));

baml_language/crates/baml_lsp2_actions_tests/src/memoization.rs

@@ -0,0 +1,108 @@
+use std::{
+    path::{Path, PathBuf},
+    sync::{Arc, Mutex},
+};
+
+use baml_lsp2_actions::{annotations, semantic_tokens};
+use baml_project::ProjectDatabase;
+
+type EventLog = Arc<Mutex<Vec<salsa::Event>>>;
+
+const SOURCE_V1: &str = r#"
+function Echo(value: int) -> int {
+    value
+}
+
+function Main() -> int {
+    let result = Echo(1)
+    result
+}
+"#;
+
+const SOURCE_V2: &str = r#"
+function Echo(value: string) -> string {
+    value
+}
+
+function Main() -> string {
+    let result = Echo("changed")
+    result
+}
+"#;
+
+fn test_db() -> (ProjectDatabase, EventLog, PathBuf) {
+    let events = EventLog::default();
+    let callback_events = Arc::clone(&events);
+    let mut db = ProjectDatabase::new_with_event_callback(Box::new(move |event| {
+        callback_events
+            .lock()
+            .expect("event log mutex poisoned")
+            .push(event);
+    }));
+    db.set_project_root(Path::new("/test"));
+    let path = PathBuf::from("/test/main.baml");
+    db.add_or_update_file(&path, SOURCE_V1);
+    (db, events, path)
+}
+
+fn clear_events(events: &EventLog) {
+    events.lock().expect("event log mutex poisoned").clear();
+}
+
+fn query_execution_count(db: &ProjectDatabase, events: &EventLog, query_name: &str) -> usize {
+    events
+        .lock()
+        .expect("event log mutex poisoned")
+        .iter()
+        .filter(|event| {
+            let salsa::EventKind::WillExecute { database_key } = &event.kind else {
+                return false;
+            };
+            let name =
+                (db as &dyn salsa::Database).ingredient_debug_name(database_key.ingredient_index());
+            name == query_name || name.ends_with(&format!("::{query_name}"))
+        })
+        .count()
+}
+
+#[test]
+fn annotations_are_memoized_and_invalidated_by_source_edits() {
+    let (mut db, events, path) = test_db();
+    let file = db.get_file(&path).expect("test file should exist");
+
+    clear_events(&events);
+    let cold = annotations(&db, file).clone();
+    assert_eq!(query_execution_count(&db, &events, "annotations"), 1);
+
+    clear_events(&events);
+    let warm = annotations(&db, file).clone();
+    assert_eq!(warm, cold);
+    assert_eq!(query_execution_count(&db, &events, "annotations"), 0);
+
+    db.add_or_update_file(&path, SOURCE_V2);
+    clear_events(&events);
+    let changed = annotations(&db, file).clone();
+    assert_ne!(changed, cold);
+    assert_eq!(query_execution_count(&db, &events, "annotations"), 1);
+}
+
+#[test]
+fn semantic_tokens_are_memoized_and_invalidated_by_source_edits() {
+    let (mut db, events, path) = test_db();
+    let file = db.get_file(&path).expect("test file should exist");
+
+    clear_events(&events);
+    let cold = semantic_tokens(&db, file).clone();
+    assert_eq!(query_execution_count(&db, &events, "semantic_tokens"), 1);
+
+    clear_events(&events);
+    let warm = semantic_tokens(&db, file).clone();
+    assert_eq!(warm, cold);
+    assert_eq!(query_execution_count(&db, &events, "semantic_tokens"), 0);
+
+    db.add_or_update_file(&path, SOURCE_V2);
+    clear_events(&events);
+    let changed = semantic_tokens(&db, file).clone();
+    assert_ne!(changed, cold);
+    assert_eq!(query_execution_count(&db, &events, "semantic_tokens"), 1);
+}

baml_language/crates/baml_lsp2_actions_tests/src/runner.rs

@@ -167,7 +167,7 @@ pub fn run_test(parsed: &ParsedTestFile) -> TestResult {
         for (filename, source_file) in &file_map {
             let hints = annotations(&db, *source_file);
             for hint in hints {
-                all_hints.push((filename.clone(), hint));
+                all_hints.push((filename.clone(), hint.clone()));
             }
         }
 
@@ -184,7 +184,7 @@ pub fn run_test(parsed: &ParsedTestFile) -> TestResult {
         for (filename, source_file) in &file_map {
             let tokens = semantic_tokens(&db, *source_file);
             for token in tokens {
-                all_tokens.push((filename.clone(), token));
+                all_tokens.push((filename.clone(), token.clone()));
             }
         }