Support multiroot workspaces (#1798)

Simplifies a bunch of logic for keeping track of baml projects
<!-- ELLIPSIS_HIDDEN -->

----

> [!IMPORTANT]
> Add support for multi-root workspaces by managing multiple `baml_src`
directories in the language server.
> 
>   - **Behavior**:
> - Support for multi-root workspaces by managing multiple `baml_src`
directories in `Session`.
>     - Removed restriction on single workspace in `server.rs`.
> - Added `get_or_create_project()` in `Session` to handle project
lookup and creation.
>   - **API**:
> - Added `getBAMLFunctions` request in `api.rs` to list functions from
all projects.
> - Updated `DidChangeTextDocumentHandler`,
`DidOpenTextDocumentHandler`, and other notification handlers to use
`get_or_create_project()`.
> - Updated `CodeLens`, `Completion`, `GotoDefinition`, `Hover`, and
`Rename` request handlers to support multi-root workspaces.
>   - **Diagnostics**:
> - Updated `publish_diagnostics()` and `project_diagnostics()` in
`diagnostics.rs` to handle diagnostics for multiple projects.
>   - **Misc**:
>     - Added tests for `get_or_create_project()` in `session.rs`.
>     - Minor logging improvements in `server.rs` and `session.rs`.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 967e24b. It will automatically
update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->

Author: aaronvg

engine/language_server/src/baml_project/file_utils.rs

@@ -19,6 +19,11 @@ use lsp_types::{TextDocumentItem, Url};
 ///   or `None` otherwise.
 pub fn find_top_level_parent(file_path: &Path) -> Option<PathBuf> {
     let mut current = file_path;
+    if let Some(file_name) = current.file_name() {
+        if file_name == "baml_src" {
+            return Some(current.to_path_buf());
+        }
+    }
     while let Some(parent) = current.parent() {
         if let Some(dir_name) = parent.file_name() {
             if dir_name == "baml_src" {

engine/language_server/src/baml_project/mod.rs

@@ -90,6 +90,15 @@ impl BamlProject {
         }
     }
 
+    pub fn list_functions(&mut self) -> Vec<BamlFunction> {
+        let runtime = self.runtime(HashMap::new());
+        if let Ok(runtime) = runtime {
+            runtime.list_functions()
+        } else {
+            vec![]
+        }
+    }
+
     pub fn check_version(
         &self,
         generator_config: &BamlGeneratorConfig,

engine/language_server/src/server.rs

@@ -1,6 +1,9 @@
 //! Scheduling, I/O, and API endpoints.
 
 use log::info;
+use lsp_types::{
+    WorkspaceClientCapabilities, WorkspaceFoldersServerCapabilities, WorkspaceServerCapabilities,
+};
 use std::num::NonZeroUsize;
 // The new PanicInfoHook name requires MSRV >= 1.82
 #[allow(deprecated)]
@@ -98,6 +101,10 @@ impl Server {
             });
             (url, settings)
         };
+        tracing::info!(
+            "--- workspace folders: {:?}",
+            init_params.workspace_folders.clone()
+        );
 
         let workspaces = init_params
             .workspace_folders
@@ -123,10 +130,6 @@ impl Server {
                 anyhow::anyhow!("Failed to get the current working directory while creating a default workspace.")
             })?;
 
-        if workspaces.len() > 1 {
-            // TODO(dhruvmanila): Support multi-root workspaces
-            anyhow::bail!("Multi-root workspaces are not supported yet");
-        }
         // for some reason tracing logs are not available before this point
         tracing::info!("Starting server with {} worker threads", worker_threads);
 
@@ -337,6 +340,13 @@ impl Server {
                     ..Default::default()
                 },
             )),
+            workspace: Some(WorkspaceServerCapabilities {
+                workspace_folders: Some(WorkspaceFoldersServerCapabilities {
+                    supported: Some(true),
+                    change_notifications: Some(lsp_types::OneOf::Left(true)),
+                }),
+                ..Default::default()
+            }),
             ..Default::default()
         }
     }

engine/language_server/src/server/api.rs

@@ -9,7 +9,7 @@ use lsp_types::{
     DidChangeTextDocumentParams, DocumentDiagnosticReport, DocumentDiagnosticReportResult,
     FullDocumentDiagnosticReport, RelatedFullDocumentDiagnosticReport,
 };
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use url::Url;
 
 mod diagnostics;
@@ -32,6 +32,19 @@ use super::{
     Result,
 };
 
+#[derive(serde::Serialize, serde::Deserialize)]
+struct BamlFunctionSpan {
+    file_path: String,
+    start: usize,
+    end: usize,
+}
+#[derive(serde::Serialize, serde::Deserialize)]
+struct BamlFunctionResult {
+    name: String,
+    span: BamlFunctionSpan,
+}
+
+struct BamlFunctionArg {}
 // --- Add debounce duration constant ---
 const DID_CHANGE_DEBOUNCE_DURATION: Duration = Duration::from_millis(250);
 
@@ -68,6 +81,51 @@ pub(super) fn request<'a>(req: lsp_server::Request) -> Task<'a> {
                 BackgroundSchedule::LatencySensitive,
             )
         }
+        "getBAMLFunctions" => {
+            tracing::info!("getBAMLFunctions");
+            return Task::local(move |session, notifier, requester, responder| {
+                let result: anyhow::Result<(serde_json::Value,)> = (|| {
+                    let mut all_functions = Vec::new();
+                    let projects = session.baml_src_projects.lock().unwrap();
+
+                    for (_, project) in projects.iter() {
+                        let functions = project
+                            .lock()
+                            .unwrap()
+                            .baml_project
+                            .list_functions()
+                            .iter()
+                            .map(|f| BamlFunctionResult {
+                                name: f.name.clone(),
+                                span: BamlFunctionSpan {
+                                    file_path: f.span.file_path.clone(),
+                                    start: f.span.start,
+                                    end: f.span.end,
+                                },
+                            })
+                            .collect::<Vec<BamlFunctionResult>>();
+
+                        all_functions.extend(functions);
+                    }
+
+                    let result = serde_json::to_value(all_functions);
+                    if let Ok(result) = result {
+                        Ok((result,))
+                    } else {
+                        Err(anyhow::anyhow!(
+                            "Failed to serialize functions: {:?}",
+                            result
+                        ))
+                    }
+                })();
+                if let Ok((result,)) = result {
+                    responder.respond(id, Ok(result)).unwrap();
+                } else {
+                    // no action
+                    // responder.respond(id, Err(result.unwrap_err())).unwrap();
+                }
+            });
+        }
         "requestDiagnostics" => {
             tracing::info!("---- requestDiagnostics");
             return Task::local(move |session, notifier, requester, responder| {
@@ -78,10 +136,8 @@ pub(super) fn request<'a>(req: lsp_server::Request) -> Task<'a> {
                         .map_err(|e| anyhow::anyhow!("Failed to parse JSON: {e}"))?;
                     let url = Url::parse(&params.project_id)
                         .map_err(|e| anyhow::anyhow!("Failed to parse URL: {e}"))?;
-                    session.ensure_project_db_for_baml_file(&url)?;
-                    // tracing::info!("url: {:?}", url);
                     let project = session
-                        .project_db_for_path(url.to_file_path().unwrap())
+                        .get_or_create_project(&url.to_file_path().unwrap())
                         .expect("Already checked for project's existence");
                     project.lock().unwrap().update_runtime(Some(notifier))?;
 
@@ -228,9 +284,9 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
         };
         info!(
             "session.projects.len(): {:?}",
-            session.projects_by_workspace_folder.lock().unwrap().len()
+            session.baml_src_projects.lock().unwrap().len()
         );
-        let _db = session.project_db_for_path(path).clone();
+        let _db = session.get_or_create_project(&path).clone();
         if _db.is_none() {
             tracing::error!("Could not find project for path");
             return Box::new(|_, _| {});

engine/language_server/src/server/api/diagnostics.rs

@@ -87,13 +87,8 @@ pub fn publish_session_lsp_diagnostics(
 ) -> Result<()> {
     // let keys = session.index().documents.keys();
     let path = file_url.to_file_path().unwrap_or(PathBuf::new());
-    let _ = session
-        .ensure_project_db_for_baml_file(file_url)
-        .map_err(|e| {
-            tracing::error!("Failed to ensure project db for baml file: {}", e);
-        });
     let project = session
-        .project_db_for_path_mut(path)
+        .get_or_create_project(&path)
         .expect("We just ensured the session is valid");
 
     let diagnostics = project_diagnostics(project.clone());

engine/language_server/src/server/api/notifications/did_change.rs

@@ -24,22 +24,24 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
         _requester: &mut Requester,
         params: DidChangeTextDocumentParams,
     ) -> Result<()> {
-        tracing::info!("------- DidChangeTextDocumentHandler");
+        tracing::info!("DidChangeTextDocumentHandler");
         let start_time_total = Instant::now();
 
         let url = params.text_document.uri;
         let path = url
             .to_file_path()
             .internal_error_msg("Could not convert URL to path")?;
-        session
-            .ensure_project_db_for_baml_file(&url)
-            .internal_error()?;
 
-        let project = session
-            .project_db_for_path_mut(path)
-            .expect("We ensured above that the project exists");
+        // Get or create the project using the unified method
+        let project = session.get_or_create_project(&path);
+        if project.is_none() {
+            tracing::error!("Failed to get or create project for path: {:?}", path);
+            show_err_msg!("Failed to get or create project for path: {:?}", path);
+        }
+
+        let project = project.unwrap();
         let document_key =
-            DocumentKey::from_url(project.lock().unwrap().root_path(), &url).internal_error()?;
+            DocumentKey::from_url(&project.lock().unwrap().root_path(), &url).internal_error()?;
 
         session
             .update_text_document(
@@ -52,11 +54,7 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
 
         tracing::info!("publishing diagnostics");
 
-        publish_diagnostics(
-            &notifier,
-            project.clone(),
-            Some(params.text_document.version),
-        )?;
+        publish_diagnostics(&notifier, project, Some(params.text_document.version))?;
 
         let elapsed = start_time_total.elapsed();
         tracing::info!("didchange total took {:?}ms", elapsed.as_millis());

engine/language_server/src/server/api/notifications/did_change_configuration.rs

@@ -17,7 +17,6 @@ impl super::SyncNotificationHandler for DidChangeConfiguration {
         _requester: &mut Requester,
         _params: types::DidChangeConfigurationParams,
     ) -> Result<()> {
-        // TODO(jane): get this wired up after the pre-release
         Ok(())
     }
 }

engine/language_server/src/server/api/notifications/did_close.rs

@@ -33,7 +33,7 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
             .to_file_path()
             .internal_error_msg("Could not convert URL to path")?;
 
-        match session.project_db_for_path(path) {
+        match session.get_or_create_project(&path) {
             None => {}
             Some(project) => {
                 let document_key = DocumentKey::from_url(

engine/language_server/src/server/api/notifications/did_open.rs

@@ -24,9 +24,14 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
         tracing::info!("DidOpenTextDocumentHandler");
 
         let url = params.text_document.uri;
-        session
-            .ensure_project_db_for_baml_file(&url)
-            .internal_error()?;
+        let file_path = url.to_file_path().internal_error()?;
+
+        let project = session.get_or_create_project(&file_path);
+        if project.is_none() {
+            tracing::error!("Failed to get or create project for path: {:?}", file_path);
+            show_err_msg!("Failed to get or create project for path: {:?}", file_path);
+        }
+
         session.reload(Some(notifier.clone())).internal_error()?;
 
         publish_session_lsp_diagnostics(&notifier, session, &url)?;

engine/language_server/src/server/api/notifications/did_save_text_document.rs

@@ -28,10 +28,7 @@ impl super::SyncNotificationHandler for DidSaveTextDocument {
         session.reload(Some(notifier.clone())).internal_error()?;
         tracing::info!("About to run generator. URL path: {:?}", path);
         session
-            .ensure_project_db_for_baml_file(&url)
-            .internal_error()?;
-        session
-            .project_db_for_path_mut(path)
+            .get_or_create_project(&path)
             .expect("Ensured that a project db exists")
             .lock()
             .unwrap()