fix(language-server): handle non-baml-src baml files

engine/language_server/src/server/api.rs

@@ -3,6 +3,7 @@ use std::{
     time::{Duration, Instant},
 };
 
+use anyhow::Context;
 use diagnostics::{file_diagnostics, project_diagnostics};
 use log::info;
 use lsp_server;
@@ -135,12 +136,12 @@ pub(super) fn request<'a>(req: lsp_server::Request) -> Task<'a> {
 
                     let params = serde_json::from_value::<DiagnosticRequestParams>(req.params)
                         .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}"))?;
-                    if !url.to_string().contains("baml_src") {
-                        return Ok(());
-                    }
+                    let url = Url::parse(&params.project_id).context("Failed to parse URL")?;
 
+                    let Ok(project) = session.get_or_create_project(url.to_file_path().unwrap())
+                    else {
+                        return Ok(());
+                    };
                     let project = session
                         .get_or_create_project(url.to_file_path().unwrap())
                         .expect("Already checked for project's existence");
@@ -299,22 +300,19 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
         .to_file_path()
         .internal_error_msg("Could not convert URL to path")?;
     Ok(Task::background(schedule, move |session: &Session| {
-        let Some(_snapshot) = session.take_snapshot(url) else {
+        let Some(snapshot) = session.take_snapshot(url) else {
             return Box::new(|_, _| {});
         };
         // info!(
         //     "session.projects.len(): {:?}",
         //     session.baml_src_projects.lock().len()
         // );
-        let _db = session.get_or_create_project(&path).clone();
-        if _db.is_none() {
-            tracing::error!("Could not find project for path");
+        let Ok(project) = session.get_or_create_project(&path) else {
             return Box::new(|_, _| {});
-        }
-        let _db = _db.unwrap();
+        };
 
-        Box::new(move |_notifier, _responder| {
-            let _ = R::run_with_snapshot(_snapshot, _db, _notifier, params);
+        Box::new(move |notifier, _responder| {
+            let _ = R::run_with_snapshot(snapshot, project, notifier, params);
         })
     }))
 }

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

@@ -95,13 +95,13 @@ pub fn publish_session_lsp_diagnostics(
 ) -> Result<()> {
     // let keys = session.index().documents.keys();
     let path = file_url.to_file_path().unwrap_or_default();
-    if !file_url.to_string().contains("baml_src") {
+    let Ok(project) = session.get_or_create_project(&path) else {
+        tracing::info!(
+            "BAML file not in baml_src directory, not publishing diagnostics: {}",
+            file_url
+        );
         return Ok(());
-    }
-    tracing::info!("publishing diagnostics for {}", file_url);
-    let project = session
-        .get_or_create_project(&path)
-        .expect("We just ensured the session is valid.");
+    };
 
     let default_flags = vec!["beta".to_string()];
     let feature_flags = session
@@ -458,3 +458,27 @@ fn ensure_absolute(project_root: &Path, file_path: &Path) -> PathBuf {
         project_root.join(file_path_relative)
     }
 }
+
+/// Creates an error diagnostic for BAML files outside baml_src directories
+pub fn not_in_baml_src_diagnostic(file_url: &Url) -> lsp_types::PublishDiagnosticsParams {
+    let range = lsp_types::Range::new(
+        lsp_types::Position::new(0, 0),
+        // Choose a position reasonably likely to be either at or past the end of the file.
+        // IDEs should correctly defend against this, ideally clamping it to the end of the file.
+        lsp_types::Position::new(10_000, 0),
+    );
+
+    lsp_types::PublishDiagnosticsParams {
+                        uri: file_url.clone(),
+                        diagnostics: vec![lsp_types::Diagnostic::new(
+            range,
+            Some(lsp_types::DiagnosticSeverity::ERROR),
+            None,
+            None,
+            "BAML files must be placed in a baml_src/ directory, see https://docs.boundaryml.com/guide/introduction/baml_src.".to_string(),
+            None,
+            None,
+        )],
+                        version: None,
+                    }
+}

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

@@ -8,7 +8,7 @@ use playground_server::WebviewRouterMessage;
 use crate::{
     server::{
         api::{
-            diagnostics::publish_diagnostics,
+            diagnostics::{not_in_baml_src_diagnostic, publish_diagnostics},
             traits::{NotificationHandler, SyncNotificationHandler},
             ResultExt,
         },
@@ -36,22 +36,19 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
         let start_time_total = Instant::now();
 
         let url = params.text_document.uri;
-        if !url.to_string().contains("baml_src") {
-            return Ok(());
-        }
-
         let path = url
             .to_file_path()
             .internal_error_msg("Could not convert URL to path")?;
 
         // 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 Ok(project) = session.get_or_create_project(&path) else {
+            notifier
+                .notify::<lsp_types::notification::PublishDiagnostics>(not_in_baml_src_diagnostic(
+                    &url,
+                ))
+                .internal_error()?;
+            return Ok(());
+        };
         let document_key =
             DocumentKey::from_url(project.lock().root_path(), &url).internal_error()?;
 

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

@@ -24,11 +24,12 @@ impl super::SyncNotificationHandler for DidChangeWatchedFiles {
         params: types::DidChangeWatchedFilesParams,
     ) -> Result<()> {
         tracing::info!("#### DidChangeWatchedFiles {:?}", params.changes);
-        if !params
-            .changes
-            .iter()
-            .any(|change| change.uri.to_string().contains("baml_src"))
-        {
+        if !params.changes.iter().any(|change| {
+            let Ok(path) = change.uri.to_file_path() else {
+                return true;
+            };
+            session.get_or_create_project(&path).is_ok()
+        }) {
             return Ok(());
         }
 

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

@@ -1,7 +1,9 @@
 use std::path::PathBuf;
 
 use lsp_server::ErrorCode;
-use lsp_types::{notification::DidCloseTextDocument, DidCloseTextDocumentParams};
+use lsp_types::{
+    notification::DidCloseTextDocument, DidCloseTextDocumentParams, PublishDiagnosticsParams,
+};
 
 // use crate::server::api::diagnostics::clear_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
@@ -27,37 +29,34 @@ impl NotificationHandler for DidCloseTextDocumentHandler {
 impl SyncNotificationHandler for DidCloseTextDocumentHandler {
     fn run(
         session: &mut Session,
-        _notifier: Notifier,
+        notifier: Notifier,
         _requester: &mut Requester,
         params: DidCloseTextDocumentParams,
     ) -> Result<()> {
         let url = params.text_document.uri;
-        if !url.to_string().contains("baml_src") {
-            return Ok(());
-        }
-
         let path = url
             .to_file_path()
             .internal_error_msg("Could not convert URL to path")?;
-
-        match session.get_or_create_project(&path) {
-            None => {}
-            Some(project) => {
-                let document_key =
-                    DocumentKey::from_url(&PathBuf::from(project.lock().root_path()), &url)
-                        .internal_error()?;
-                session
-                    .close_document(&document_key)
-                    .with_failure_code(ErrorCode::InternalError)?;
-                // Remove the unsaved file from the project as well
-                // TODO: ideally the baml project just has a view of unsaved files directly from the Session itself, and not maintain its own state / copy of the unsaved files
-                project
-                    .lock()
-                    .baml_project
-                    .remove_unsaved_file(&document_key);
-            }
+        let Ok(project) = session.get_or_create_project(&path) else {
+            return Ok(());
+        };
+
+        let document_key = DocumentKey::from_url(&PathBuf::from(project.lock().root_path()), &url)
+            .internal_error()?;
+        session
+            .close_document(&document_key)
+            .with_failure_code(ErrorCode::InternalError)?;
+        // Remove the unsaved file from the project as well
+        // TODO: ideally the baml project just has a view of unsaved files directly from the Session itself, and not maintain its own state / copy of the unsaved files
+        {
+            // nested to block to ensure the project is unlocked before session.reload()
+            project
+                .lock()
+                .baml_project
+                .remove_unsaved_file(&document_key);
         }
-        session.reload(Some(_notifier)).internal_error()?;
+
+        session.reload(Some(notifier)).internal_error()?;
 
         Ok(())
     }

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

@@ -6,7 +6,7 @@ use lsp_types::{
 use crate::{
     server::{
         api::{
-            diagnostics::publish_session_lsp_diagnostics,
+            diagnostics::{not_in_baml_src_diagnostic, publish_session_lsp_diagnostics},
             notifications::{
                 baml_src_version::BamlSrcVersionPayload,
                 did_save_text_document::send_generator_version,
@@ -41,9 +41,17 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
         tracing::info!("DidOpenTextDocumentHandler");
 
         let url = params.text_document.uri;
-        if !url.to_string().contains("baml_src") {
+        let path = url
+            .to_file_path()
+            .internal_error_msg("Could not convert URL to path")?;
+        let Ok(project) = session.get_or_create_project(&path) else {
+            notifier
+                .notify::<lsp_types::notification::PublishDiagnostics>(not_in_baml_src_diagnostic(
+                    &url,
+                ))
+                .internal_error()?;
             return Ok(());
-        }
+        };
 
         // TODO: do this when server initializes instead of every time a file is opened
         // note this just schedules the task. It will run after the current task is done.
@@ -67,12 +75,8 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
             )
             .internal_error()?;
 
-        let file_path = url
-            .to_file_path()
-            .internal_error_msg(&format!("Could not convert URL '{url}' to file path"))?;
-
         // tracing::info!("before get_or_create_project");
-        if let Some(project) = session.get_or_create_project(&file_path) {
+        {
             let locked = project.lock();
             let default_flags = vec!["beta".to_string()];
             let effective_flags = session
@@ -84,11 +88,7 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
 
             let generator_version = locked.get_common_generator_version();
             send_generator_version(&notifier, &locked, generator_version.as_ref().ok());
-        } else {
-            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);
         }
-        tracing::info!("after get_or_create_project");
 
         // session.open_text_document(
         //     DocumentKey::from_path(&file_path, &file_path).internal_error()?,

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

@@ -1,11 +1,17 @@
 use std::borrow::Cow;
 
-use lsp_types::{self as types, notification as notif, request::Request, ConfigurationParams};
+use lsp_types::{
+    self as types, notification as notif, request::Request, ConfigurationParams,
+    PublishDiagnosticsParams,
+};
 
 use crate::{
     baml_project::{common_version_up_to_patch, Project},
     server::{
-        api::{self, notifications::baml_src_version::BamlSrcVersionPayload, ResultExt},
+        api::{
+            self, diagnostics::not_in_baml_src_diagnostic,
+            notifications::baml_src_version::BamlSrcVersionPayload, ResultExt,
+        },
         client::{Notifier, Requester},
         Result, Task,
     },
@@ -27,13 +33,18 @@ impl super::SyncNotificationHandler for DidSaveTextDocument {
     ) -> Result<()> {
         tracing::info!("Did save text document---------");
         let url = params.text_document.uri;
-        if !url.to_string().contains("baml_src") {
-            return Ok(());
-        }
-
         let path = url
             .to_file_path()
             .internal_error_msg("Could not convert URL to path")?;
+        let Ok(project) = session.get_or_create_project(&path) else {
+            notifier
+                .notify::<lsp_types::notification::PublishDiagnostics>(not_in_baml_src_diagnostic(
+                    &url,
+                ))
+                .internal_error()?;
+            return Ok(());
+        };
+
         session.clear_unsaved_files();
         session.reload(Some(notifier.clone())).internal_error()?;
 
@@ -45,9 +56,6 @@ impl super::SyncNotificationHandler for DidSaveTextDocument {
         }
 
         tracing::info!("About to run generator. URL path: {:?}", path);
-        let project = session
-            .get_or_create_project(&path)
-            .expect("Ensured that a project db exists");
         let mut locked = project.lock();
 
         let default_flags = vec!["beta".to_string()];
@@ -164,30 +172,32 @@ impl super::BackgroundDocumentNotificationHandler for DidSaveTextDocument {
 
         // Note: In the background version, we need to get the project from the snapshot
         // instead of modifying the session directly
-        if let Some(project) = snapshot.project() {
-            let default_flags = vec!["beta".to_string()];
-            let effective_flags = snapshot
-                .session_baml_settings()
-                .feature_flags
-                .as_ref()
-                .unwrap_or(&default_flags);
-            project.lock().run_generators_without_debounce(
-                effective_flags,
-                |message| {
-                    tracing::info!("About to notify client that generator has run.");
-                    notifier.notify_baml_info(&message).unwrap_or(())
-                },
-                |e| {
-                    tracing::error!("Error generating: {e}");
-                    notifier.notify_baml_error(&e).unwrap_or(())
-                },
-            );
-        } else {
+
+        let Ok(project) = snapshot.project() else {
             tracing::error!("No project found in snapshot for file {:?}", path);
             notifier
                 .notify_baml_error(&format!("No project found for file {path:?}"))
                 .unwrap_or(());
-        }
+            return Ok(());
+        };
+
+        let default_flags = vec!["beta".to_string()];
+        let effective_flags = snapshot
+            .session_baml_settings()
+            .feature_flags
+            .as_ref()
+            .unwrap_or(&default_flags);
+        project.lock().run_generators_without_debounce(
+            effective_flags,
+            |message| {
+                tracing::info!("About to notify client that generator has run.");
+                notifier.notify_baml_info(&message).unwrap_or(())
+            },
+            |e| {
+                tracing::error!("Error generating: {e}");
+                notifier.notify_baml_error(&e).unwrap_or(())
+            },
+        );
 
         Ok(())
     }