feat(lsp): resolution-first semantic tokens rework

baml_language/Cargo.toml

@@ -81,6 +81,7 @@ sys_native = { path = "crates/sys_native", default-features = false }
 baml_db = { path = "crates/baml_db" }
 baml_fmt = { path = "crates/baml_fmt" }
 baml_lsp2_actions = { path = "crates/baml_lsp2_actions" }
+baml_lsp2_actions_tests = { path = "crates/baml_lsp2_actions_tests" }
 baml_lsp_server = { path = "crates/baml_lsp_server", default-features = false }
 baml_project = { path = "crates/baml_project" }
 baml_release = { path = "crates/baml_release" }
@@ -128,6 +129,7 @@ aws-sigv4 = { path = "forks/aws-sigv4", package = "forked_aws_sigv4" }
 aws-bedrock = { path = "forks/aws-bedrock", package = "forked_aws_bedrock" }
 google-cloud-auth = { path = "forks/google-cloud-auth", package = "forked_google_cloud_auth" }
 axum = { version = "0.8.4", default-features = false, features = [ "http1", "matched-path", "original-uri", "query", "tokio", "ws" ] }
+bitflags = "2"
 # Native `borsh` derives on the entire Program type graph (bex_vm_types +
 # baml_type + baml_exec envelope). The pack envelope was previously
 # bitcode-with-serde; borsh's native derive avoids both the bitcode crate

baml_language/crates/baml_cli/src/paint.rs

@@ -103,7 +103,8 @@ fn style_for(token_type: SemanticTokenType) -> Style {
         T::Parameter => Style::new().color256(173),
         T::Namespace => Style::new().cyan().bright(),
         T::String | T::Regexp => Style::new().green(),
-        T::Number => Style::new().yellow().bright(),
+        T::EscapeSequence => Style::new().color256(214),
+        T::Number | T::Boolean => Style::new().yellow().bright(),
         T::Comment => Style::new().color256(244),
         T::Decorator => Style::new().color256(179),
         T::Operator => Style::new().color256(245),

baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs

@@ -52,6 +52,14 @@ fn is_ident_token(kind: SyntaxKind) -> bool {
             | SyntaxKind::KW_INTERFACE
             | SyntaxKind::KW_EXTENDS
             | SyntaxKind::KW_REQUIRES
+            // Contextual keywords re-lexed from a `Word`: still lower by text
+            // (the literal/identifier arms below switch on the text), so they
+            // must read as ident tokens just as they did when they were `Word`.
+            | SyntaxKind::KW_AS
+            | SyntaxKind::KW_TYPE
+            | SyntaxKind::KW_TRUE
+            | SyntaxKind::KW_FALSE
+            | SyntaxKind::KW_NULL
     )
 }
 

baml_language/crates/baml_compiler2_hir/src/builder.rs

@@ -960,7 +960,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         let binding_visible_from = source_map.pattern_span(clause.binding).start();
         self.register_local_pattern(
             clause.binding,
-            DefinitionSite::PatternBinding(clause.binding),
+            DefinitionSite::CatchBinding(clause.binding),
             body,
             source_map,
             binding_visible_from,
@@ -969,7 +969,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             let st_visible_from = source_map.pattern_span(st_pat).start();
             self.register_local_pattern(
                 st_pat,
-                DefinitionSite::PatternBinding(st_pat),
+                DefinitionSite::CatchBinding(st_pat),
                 body,
                 source_map,
                 st_visible_from,
@@ -1436,6 +1436,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         self.class_depth -= 1;
 
         let local_id = self.item_tree.alloc_interface(i, default_method_ids);
+        ItemTree::collect_interface_spans(&mut self.item_tree_source_map, local_id, i);
         let loc = InterfaceLoc::new(self.db, self.file, local_id);
         self.type_contributions.push((
             i.name.clone(),

baml_language/crates/baml_compiler2_hir/src/item_tree.rs

@@ -408,6 +408,11 @@ pub struct ItemTreeSourceMap {
     pub enum_variant_spans: FxHashMap<LocalItemId<EnumMarker>, Vec<TextRange>>,
     /// `name_span` for each function.
     pub function_name_spans: FxHashMap<LocalItemId<FunctionMarker>, TextRange>,
+    /// `name_span` for each interface's fields, parallel to `Interface::fields`.
+    pub interface_field_spans: FxHashMap<LocalItemId<InterfaceMarker>, Vec<TextRange>>,
+    /// `name_span` for each interface's required methods, parallel to
+    /// `Interface::required_methods`.
+    pub interface_method_spans: FxHashMap<LocalItemId<InterfaceMarker>, Vec<TextRange>>,
 }
 
 // ── ItemTree ─────────────────────────────────────────────────────────────────
@@ -693,6 +698,22 @@ impl ItemTree {
         source_map.enum_variant_spans.insert(id, spans);
     }
 
+    /// Populate source map spans for an interface allocated via `alloc_interface`.
+    pub fn collect_interface_spans(
+        source_map: &mut ItemTreeSourceMap,
+        id: LocalItemId<InterfaceMarker>,
+        iface_def: &ast::InterfaceDef,
+    ) {
+        let field_spans: Vec<TextRange> = iface_def.fields.iter().map(|f| f.name_span).collect();
+        source_map.interface_field_spans.insert(id, field_spans);
+        let method_spans: Vec<TextRange> = iface_def
+            .required_methods
+            .iter()
+            .map(|m| m.name_span)
+            .collect();
+        source_map.interface_method_spans.insert(id, method_spans);
+    }
+
     /// Allocate an interface (BEP-044) in the `ItemTree`.
     ///
     /// `default_method_ids` are the `FunctionMarker` ids for any default

baml_language/crates/baml_compiler2_hir/src/semantic_index.rs

@@ -50,8 +50,12 @@ pub enum DefinitionSite {
     Statement(StmtId),
     /// Defined as a function parameter (with its index).
     Parameter(usize),
-    /// Defined by a pattern binding (match arm, catch arm, catch clause, etc.).
+    /// Defined by a pattern binding (match arm, if-let, etc.).
     PatternBinding(PatId),
+    /// The error (and optional stack-trace) binding of a `catch (e) { … }`
+    /// clause. Like a function parameter — a value bound by the clause and
+    /// scoped to its body — so it is highlighted as a parameter.
+    CatchBinding(PatId),
 }
 
 // ── BindingId ────────────────────────────────────────────────────────────────

baml_language/crates/baml_compiler2_mir/src/lower.rs

@@ -1184,7 +1184,10 @@ fn resolution_to_item_ref(
                 name: func_data.name.clone(),
             })
         }
-        MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+        MemberResolution::Field { .. }
+        | MemberResolution::Variant { .. }
+        | MemberResolution::InterfaceMethod { .. }
+        | MemberResolution::InterfaceField { .. } => None,
     }
 }
 
@@ -6089,6 +6092,13 @@ impl<'db> LoweringContext<'db> {
                     Some(MemberResolution::Variant { .. }) => {
                         // Handled by expr_types check below.
                     }
+                    Some(
+                        MemberResolution::InterfaceMethod { .. }
+                        | MemberResolution::InterfaceField { .. },
+                    ) => {
+                        // Interface member on an interface-typed value: fall through
+                        // to the runtime interface-dispatch path below.
+                    }
                     None => {}
                 }
             }
@@ -6157,6 +6167,11 @@ impl<'db> LoweringContext<'db> {
                         self.lower_multi_segment_path_as_field_chain(expr_id, segments, dest);
                         return;
                     }
+                    MemberResolution::InterfaceMethod { .. }
+                    | MemberResolution::InterfaceField { .. } => {
+                        // Interface member on an interface-typed value: fall through
+                        // to the runtime interface-dispatch path below.
+                    }
                 }
             }
             // An interface method referenced as a *value* on a generic- or
@@ -8676,7 +8691,10 @@ impl<'db> LoweringContext<'db> {
                         | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                             Some(*func_loc)
                         }
-                        MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                        MemberResolution::Field { .. }
+                        | MemberResolution::Variant { .. }
+                        | MemberResolution::InterfaceMethod { .. }
+                        | MemberResolution::InterfaceField { .. } => None,
                     });
                 if from_pmr.is_some() {
                     from_pmr
@@ -8690,9 +8708,10 @@ impl<'db> LoweringContext<'db> {
                             | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                                 Some(*func_loc)
                             }
-                            MemberResolution::Field { .. } | MemberResolution::Variant { .. } => {
-                                None
-                            }
+                            MemberResolution::Field { .. }
+                            | MemberResolution::Variant { .. }
+                            | MemberResolution::InterfaceMethod { .. }
+                            | MemberResolution::InterfaceField { .. } => None,
                         })
                 }
             };
@@ -8713,7 +8732,10 @@ impl<'db> LoweringContext<'db> {
                     | MemberResolution::UnboundMethod { func_loc, .. }
                     | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => Some(*func_loc),
                     MemberResolution::Free { func_loc } => Some(*func_loc),
-                    MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                    MemberResolution::Field { .. }
+                    | MemberResolution::Variant { .. }
+                    | MemberResolution::InterfaceMethod { .. }
+                    | MemberResolution::InterfaceField { .. } => None,
                 };
                 if let Some(fl) = func_loc {
                     let body = baml_compiler2_ppir::function_body(self.db, fl);
@@ -8771,7 +8793,10 @@ impl<'db> LoweringContext<'db> {
                         | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                             Some(*func_loc)
                         }
-                        MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                        MemberResolution::Field { .. }
+                        | MemberResolution::Variant { .. }
+                        | MemberResolution::InterfaceMethod { .. }
+                        | MemberResolution::InterfaceField { .. } => None,
                     });
                 if from_pmr.is_some() {
                     from_pmr
@@ -8785,9 +8810,10 @@ impl<'db> LoweringContext<'db> {
                             | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                                 Some(*func_loc)
                             }
-                            MemberResolution::Field { .. } | MemberResolution::Variant { .. } => {
-                                None
-                            }
+                            MemberResolution::Field { .. }
+                            | MemberResolution::Variant { .. }
+                            | MemberResolution::InterfaceMethod { .. }
+                            | MemberResolution::InterfaceField { .. } => None,
                         })
                 }
             };
@@ -8808,7 +8834,10 @@ impl<'db> LoweringContext<'db> {
                     | MemberResolution::UnboundMethod { func_loc, .. }
                     | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => Some(*func_loc),
                     MemberResolution::Free { func_loc } => Some(*func_loc),
-                    MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                    MemberResolution::Field { .. }
+                    | MemberResolution::Variant { .. }
+                    | MemberResolution::InterfaceMethod { .. }
+                    | MemberResolution::InterfaceField { .. } => None,
                 };
                 if let Some(fl) = func_loc {
                     let body = baml_compiler2_ppir::function_body(self.db, fl);
@@ -8861,7 +8890,10 @@ impl<'db> LoweringContext<'db> {
                         | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                             Some(*func_loc)
                         }
-                        MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                        MemberResolution::Field { .. }
+                        | MemberResolution::Variant { .. }
+                        | MemberResolution::InterfaceMethod { .. }
+                        | MemberResolution::InterfaceField { .. } => None,
                     });
                 if from_pmr.is_some() {
                     from_pmr
@@ -8875,9 +8907,10 @@ impl<'db> LoweringContext<'db> {
                             | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                                 Some(*func_loc)
                             }
-                            MemberResolution::Field { .. } | MemberResolution::Variant { .. } => {
-                                None
-                            }
+                            MemberResolution::Field { .. }
+                            | MemberResolution::Variant { .. }
+                            | MemberResolution::InterfaceMethod { .. }
+                            | MemberResolution::InterfaceField { .. } => None,
                         })
                 }
             };
@@ -8898,7 +8931,10 @@ impl<'db> LoweringContext<'db> {
                     | MemberResolution::UnboundMethod { func_loc, .. }
                     | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => Some(*func_loc),
                     MemberResolution::Free { func_loc } => Some(*func_loc),
-                    MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                    MemberResolution::Field { .. }
+                    | MemberResolution::Variant { .. }
+                    | MemberResolution::InterfaceMethod { .. }
+                    | MemberResolution::InterfaceField { .. } => None,
                 };
                 if let Some(fl) = func_loc {
                     let body = baml_compiler2_ppir::function_body(self.db, fl);
@@ -8972,7 +9008,10 @@ impl<'db> LoweringContext<'db> {
                         | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                             Some(*func_loc)
                         }
-                        MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                        MemberResolution::Field { .. }
+                        | MemberResolution::Variant { .. }
+                        | MemberResolution::InterfaceMethod { .. }
+                        | MemberResolution::InterfaceField { .. } => None,
                     });
                 if from_pmr.is_some() {
                     from_pmr
@@ -8986,9 +9025,10 @@ impl<'db> LoweringContext<'db> {
                             | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                                 Some(*func_loc)
                             }
-                            MemberResolution::Field { .. } | MemberResolution::Variant { .. } => {
-                                None
-                            }
+                            MemberResolution::Field { .. }
+                            | MemberResolution::Variant { .. }
+                            | MemberResolution::InterfaceMethod { .. }
+                            | MemberResolution::InterfaceField { .. } => None,
                         })
                 }
             };
@@ -9071,7 +9111,10 @@ impl LoweringContext<'_> {
                         | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                             Some(*func_loc)
                         }
-                        MemberResolution::Field { .. } | MemberResolution::Variant { .. } => None,
+                        MemberResolution::Field { .. }
+                        | MemberResolution::Variant { .. }
+                        | MemberResolution::InterfaceMethod { .. }
+                        | MemberResolution::InterfaceField { .. } => None,
                     });
                 if from_pmr.is_some() {
                     from_pmr
@@ -9085,9 +9128,10 @@ impl LoweringContext<'_> {
                             | MemberResolution::InterfaceDefaultMethod { func_loc, .. } => {
                                 Some(*func_loc)
                             }
-                            MemberResolution::Field { .. } | MemberResolution::Variant { .. } => {
-                                None
-                            }
+                            MemberResolution::Field { .. }
+                            | MemberResolution::Variant { .. }
+                            | MemberResolution::InterfaceMethod { .. }
+                            | MemberResolution::InterfaceField { .. } => None,
                         })
                 }
             }
@@ -9915,7 +9959,10 @@ impl<'db> LoweringContext<'db> {
                         return;
                     }
                 }
-                MemberResolution::Field { .. } | MemberResolution::Variant { .. } => {
+                MemberResolution::Field { .. }
+                | MemberResolution::Variant { .. }
+                | MemberResolution::InterfaceMethod { .. }
+                | MemberResolution::InterfaceField { .. } => {
                     // Fall through — handled by the existing field/enum-variant lowering below
                 }
             }