[LSP] Fix issue where vscode-generated baml_client would have stale data (#1791)

<!-- ELLIPSIS_HIDDEN -->



> [!IMPORTANT]
> Fixes stale data issue in `baml_client` by adjusting notification
handling and session management.
> 
>   - **Behavior**:
>     - Replaces `ruff` with `baml` in `logging.rs` for log filtering.
> - Modifies `DidSaveTextDocument` handling in `api.rs` to use
`local_notification_task` instead of `background_notification_task` to
prevent stale data.
> - Adds `clear_unsaved_files()` in `session.rs` to clear unsaved files
on save.
>   - **Session Management**:
> - Adds `clear_unsaved_files()` method in `Session` to clear unsaved
files.
> - Updates `reload()` in `Session` to handle unsaved files correctly.
>   - **Connection Handling**:
> - Changes `handle_shutdown()` in `connection.rs` to loop until an exit
notification is received, handling unexpected messages.
>   - **Client Readiness**:
> - Adds `clientReady` flag in `index.ts` to ensure client is ready
before sending requests.
> 
> <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 30c980b. It will automatically
update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->

Author: aaronvg

engine/language_server/src/logging.rs

@@ -111,7 +111,7 @@ impl<S> tracing_subscriber::layer::Filter<S> for LogLevelFilter {
         meta: &tracing::Metadata<'_>,
         _: &tracing_subscriber::layer::Context<'_, S>,
     ) -> bool {
-        let filter = if meta.target().starts_with("ruff") {
+        let filter = if meta.target().starts_with("baml") {
             self.filter.trace_level()
         } else {
             tracing::Level::INFO

engine/language_server/src/server/api.rs

@@ -83,7 +83,7 @@ pub(super) fn request<'a>(req: lsp_server::Request) -> Task<'a> {
                     let project = session
                         .project_db_for_path(url.to_file_path().unwrap())
                         .expect("Already checked for project's existence");
-                    project.lock().unwrap().update_runtime(Some(notifier));
+                    project.lock().unwrap().update_runtime(Some(notifier))?;
 
                     // TODO: I think we need to send ALL diagnostics for the project. Not sure how this report is different vs sending a signle diagnostic param message
                     let diagnostics = file_diagnostics(project.clone(), &url);
@@ -183,11 +183,15 @@ pub(super) fn notification<'a>(notif: lsp_server::Notification) -> Vec<Task<'a>>
         }
         // --- DidSaveTextDocument now uses the simple local task helper ---
         notification::DidSaveTextDocument::METHOD => {
+            tracing::info!("Did save text document---------");
             handle_notification_result_error::<notification::DidSaveTextDocument>(
-                background_notification_task::<notification::DidSaveTextDocument>(
-                    notif,
-                    BackgroundSchedule::LatencySensitive,
-                ),
+                // Do not use background notifs yet, as baml_client may not have an updated view of the project files
+                // See the did_save_text_document.rs file for more details
+                // background_notification_task::<notification::DidSaveTextDocument>(
+                //     notif,
+                //     BackgroundSchedule::LatencySensitive,
+                // ),
+                local_notification_task::<notification::DidSaveTextDocument>(notif),
             )
         }
 

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

@@ -19,10 +19,12 @@ impl super::SyncNotificationHandler for DidSaveTextDocument {
         _requester: &mut Requester,
         params: types::DidSaveTextDocumentParams,
     ) -> Result<()> {
+        tracing::info!("Did save text document---------");
         let url = params.text_document.uri;
         let path = url
             .to_file_path()
             .internal_error_msg("Could not convert URL to path")?;
+        session.clear_unsaved_files();
         session.reload(Some(notifier.clone())).internal_error()?;
         tracing::info!("About to run generator. URL path: {:?}", path);
         session
@@ -51,6 +53,8 @@ impl super::SyncNotificationHandler for DidSaveTextDocument {
     }
 }
 
+// Do not use this yet, it seems it has an outdated view of the project files and it generates
+// stale baml clients
 impl super::BackgroundDocumentNotificationHandler for DidSaveTextDocument {
     fn document_url(params: &types::DidSaveTextDocumentParams) -> Cow<types::Url> {
         Cow::Borrowed(&params.text_document.uri)

engine/language_server/src/server/connection.rs

@@ -102,19 +102,40 @@ impl Connection {
                 self.sender
                     .send(lsp::Response::new_ok(id.clone(), ()).into())?;
                 tracing::info!("Shutdown request received. Waiting for an exit notification...");
-                match self.receiver.recv_timeout(std::time::Duration::from_secs(30))? {
-                    lsp::Message::Notification(lsp::Notification { method, .. }) if method == lsp_types::notification::Exit::METHOD => {
-                        tracing::info!("Exit notification received. Server shutting down...");
-                        Ok(true)
-                    },
-                    message => anyhow::bail!("Server received unexpected message {message:?} while waiting for exit notification")
+
+                loop {
+                    match &self
+                        .receiver
+                        .recv_timeout(std::time::Duration::from_secs(30))?
+                    {
+                        lsp::Message::Notification(lsp::Notification { method, .. })
+                            if method == lsp_types::notification::Exit::METHOD =>
+                        {
+                            tracing::info!("Exit notification received. Server shutting down...");
+                            return Ok(true);
+                        }
+                        lsp::Message::Request(lsp::Request { id, method, .. }) => {
+                            tracing::warn!(
+                                "Server received unexpected request {method} ({id}) while waiting for exit notification",
+                            );
+                            self.sender.send(lsp::Message::Response(lsp::Response::new_err(
+                                id.clone(),
+                                lsp::ErrorCode::InvalidRequest as i32,
+                                "Server received unexpected request while waiting for exit notification".to_string(),
+                            )))?;
+                        }
+                        message => {
+                            tracing::warn!(
+                                "Server received unexpected message while waiting for exit notification: {message:?}"
+                            );
+                        }
+                    }
                 }
             }
             lsp::Message::Notification(lsp::Notification { method, .. })
                 if method == lsp_types::notification::Exit::METHOD =>
             {
-                tracing::error!("Server received an exit notification before a shutdown request was sent. Exiting...");
-                Ok(true)
+                anyhow::bail!("Server received an exit notification before a shutdown request was sent. Exiting...");
             }
             _ => Ok(false),
         }

engine/language_server/src/session.rs

@@ -197,6 +197,12 @@ impl Session {
 
         Ok(())
     }
+    pub fn clear_unsaved_files(&mut self) {
+        tracing::info!("Clearing unsaved files");
+        for (_folder, project) in self.projects_by_workspace_folder.lock().unwrap().iter_mut() {
+            project.lock().unwrap().baml_project.unsaved_files.clear();
+        }
+    }
 
     /// Creates a document snapshot with the URL referencing the document to snapshot.
     pub fn take_snapshot(&self, url: Url) -> Option<DocumentSnapshot> {

typescript/vscode-ext/packages/vscode/src/plugins/language-server/index.ts

@@ -27,6 +27,7 @@ import { bamlConfig, getConfig } from './bamlConfig'
 
 export { bamlConfig }
 const packageJson = require('../../../../package.json') // eslint-disable-line
+let clientReady = false
 
 let client: LanguageClient
 let serverModule: string
@@ -46,6 +47,10 @@ export const requestDiagnostics = async () => {
   if (!currentFile.endsWith('.baml')) {
     return
   }
+  if (!clientReady) {
+    console.warn('client not ready')
+    return
+  }
   await client?.sendRequest('requestDiagnostics', { projectId: currentFile })
 }
 
@@ -141,6 +146,7 @@ const activateClient = (
     .onReady()
     .then(() => {
       console.log('client ready')
+      clientReady = true
       client.createDefaultErrorHandler(2)
       requestDiagnostics()
       client.onNotification('baml/showLanguageServerOutput', () => {