feat: implement format-on-save in vscode and baml-cli fmt

engine/baml-schema-wasm/src/lib.rs

@@ -1,6 +1,7 @@
 #[cfg(target_arch = "wasm32")]
 pub mod runtime_wasm;
 
+use internal_baml_core::internal_baml_schema_ast::{format_schema, FormatOptions};
 use std::env;
 use wasm_bindgen::prelude::*;
 
@@ -9,3 +10,24 @@ pub fn version() -> String {
     // register_panic_hook();
     env!("CARGO_PKG_VERSION").to_string()
 }
+
+#[wasm_bindgen]
+pub fn format_document(path: String, text: String) -> Option<String> {
+    log::info!("Trying to format document (rust): {}", path);
+    match format_schema(
+        &text,
+        FormatOptions {
+            indent_width: 2,
+            fail_on_unhandled_rule: false,
+        },
+    ) {
+        Ok(formatted) => {
+            log::info!("Formatted document: {}", formatted);
+            Some(formatted)
+        }
+        Err(e) => {
+            log::error!("Failed to format document: {} {:?}", path, e);
+            None
+        }
+    }
+}

engine/cli/src/commands.rs

@@ -68,7 +68,12 @@ impl RuntimeCli {
                 args.from = BamlRuntime::parse_baml_src_path(&args.from)?;
                 t.block_on(async { args.run_async().await })
             }
-            Commands::Format(args) => args.run(),
+            Commands::Format(args) => {
+                // We deliberately don't apply parse_baml_src_path here
+                // see format.rs for more details
+                // args.from = BamlRuntime::parse_baml_src_path(&args.from)?;
+                args.run()
+            }
         }
     }
 }

engine/cli/src/format.rs

@@ -1,31 +1,64 @@
 use std::{fs, path::PathBuf};
 
 use anyhow::Result;
+use baml_runtime::{baml_src_files, BamlRuntime};
 use clap::Args;
 use internal_baml_core::internal_baml_schema_ast::{format_schema, FormatOptions};
 
 #[derive(Args, Debug)]
 pub struct FormatArgs {
     #[arg(long, help = "path/to/baml_src", default_value = "./baml_src")]
     pub from: PathBuf,
+
+    #[arg(
+        help = "Specific files to format. If none provided, formats all files in the baml_src directory"
+    )]
+    pub paths: Vec<PathBuf>,
+
+    #[arg(
+        short = 'n',
+        long = "dry-run",
+        help = "Write formatter changes to stdout instead of files",
+        default_value = "false"
+    )]
+    pub dry_run: bool,
 }
 
 impl FormatArgs {
     pub fn run(&self) -> Result<()> {
-        let source = fs::read_to_string(&self.from)?;
-        let formatted = format_schema(
-            &source,
-            FormatOptions {
-                indent_width: 4,
-                fail_on_unhandled_rule: false,
-            },
-        )?;
-
-        let mut to = self.from.clone();
-        to.set_extension("formatted.baml");
-        fs::write(&to, formatted)?;
-
-        log::info!("Formatted {} to {}", self.from.display(), to.display());
+        let paths = if self.paths.is_empty() {
+            // Usually this is done in commands.rs, but fmt is a special case
+            // because it doesn't need to actually load the BAML runtime to parse
+            // BAML files.
+            let from = BamlRuntime::parse_baml_src_path(&self.from)?;
+            baml_src_files(&from)?
+        } else {
+            self.paths.clone()
+        };
+
+        for path in paths.iter() {
+            let source = fs::read_to_string(&path)?;
+            match format_schema(
+                &source,
+                FormatOptions {
+                    indent_width: 2,
+                    fail_on_unhandled_rule: false,
+                },
+            ) {
+                Ok(formatted) => {
+                    if self.dry_run {
+                        println!("{}", formatted);
+                    } else {
+                        fs::write(&path, formatted)?;
+                    }
+                }
+                Err(e) => {
+                    log::error!("Failed to format {}: {}", path.display(), e);
+                }
+            }
+        }
+
+        log::info!("Formatted {} files", paths.len());
 
         Ok(())
     }

tools/build

@@ -137,7 +137,7 @@ case "$_path" in
 
     if [ "$_watch_mode" -eq 1 ]; then
       npx nodemon \
-        --ext rs,hb,hbs,j2,toml,baml \
+        --ext rs,hb,hbs,j2,toml \
         --watch "${_repo_root}/engine" \
         --ignore 'target' \
         --exec "${command}"
@@ -157,7 +157,7 @@ case "$_path" in
 
     if [ "$_watch_mode" -eq 1 ]; then
       npx nodemon \
-        --ext rs,hb,hbs,j2,toml,baml \
+        --ext rs,hb,hbs,j2,toml \
         --watch "${_repo_root}/engine" \
         --ignore 'target/**' \
         --exec "${command}"

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

@@ -140,7 +140,7 @@ export function startServer(options?: LSOptions): void {
       capabilities: {
         textDocumentSync: TextDocumentSyncKind.Full,
         definitionProvider: true,
-        documentFormattingProvider: false,
+        documentFormattingProvider: true,
         completionProvider: {
           resolveProvider: false,
           triggerCharacters: ['@', '"', '.'],
@@ -527,6 +527,21 @@ export function startServer(options?: LSOptions): void {
     }
   })
 
+  connection.onDocumentFormatting((params: DocumentFormattingParams) => {
+    try {
+      const doc = getDocument(params.textDocument.uri)
+      if (doc) {
+        const formatted = BamlWasm.format_document(doc.uri, doc.getText())
+        if (formatted) {
+          return [TextEdit.replace(Range.create(doc.positionAt(0), doc.positionAt(doc.getText().length)), formatted)]
+        }
+      }
+      return []
+    } catch (e) {
+      console.error(`Error occurred while formatting document:\n${e}`)
+    }
+  })
+
   connection.onCodeLens((params: CodeLensParams) => {
     try {
       const document = getDocument(params.textDocument.uri)

typescript/vscode-ext/packages/package.json

@@ -154,6 +154,11 @@
         "command": "baml.selectTestCase",
         "title": "Select Test Case",
         "category": "Baml"
+      },
+      {
+        "command": "baml.setDefaultFormatter",
+        "title": "Set Default Formatter",
+        "category": "Baml"
       }
     ]
   },

typescript/vscode-ext/packages/vscode/src/extension.ts

@@ -354,6 +354,12 @@ export function activate(context: vscode.ExtensionContext) {
     }
   })
 
+  const config = vscode.workspace.getConfiguration('editor', { languageId: 'baml' })
+  if (!config.get('defaultFormatter')) {
+    // TODO: once the BAML formatter is stable, we should auto-prompt people to set it as the default formatter.
+    // void vscode.commands.executeCommand('baml.setDefaultFormatter')
+  }
+
   // Listen for messages from the webview
 
   plugins.map(async (plugin) => {
@@ -394,7 +400,6 @@ export function deactivate(): void {
   }
   server?.close()
 }
-
 class DiagnosticCodeActionProvider implements vscode.CodeActionProvider {
   public provideCodeActions(
     document: vscode.TextDocument,

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

@@ -382,6 +382,49 @@ const plugin: BamlVSCodePlugin = {
           }
         },
       ),
+
+      commands.registerCommand('baml.setDefaultFormatter', async () => {
+        enum AutoFormatChoice {
+          Yes = 'Yes (always)',
+          OnlyInWorkspace = 'Yes (in workspace)',
+          No = 'No',
+        }
+        const selection = await vscode.window.showInformationMessage(
+          'Would you like to auto-format BAML files on save?',
+          AutoFormatChoice.Yes,
+          AutoFormatChoice.OnlyInWorkspace,
+          AutoFormatChoice.No,
+        )
+        if (selection === AutoFormatChoice.No) {
+          return
+        }
+
+        const config = vscode.workspace.getConfiguration('editor', { languageId: 'baml' })
+
+        const configTarget =
+          selection === AutoFormatChoice.Yes ? vscode.ConfigurationTarget.Global : vscode.ConfigurationTarget.Workspace
+        const overrideInLanguage = true
+
+        for (const [key, value] of Object.entries({
+          defaultFormatter: 'Boundary.baml-extension',
+          formatOnSave: true,
+        })) {
+          await config.update(key, value, configTarget, overrideInLanguage)
+        }
+
+        switch (selection) {
+          case AutoFormatChoice.Yes:
+            vscode.window.showInformationMessage(
+              'BAML files will now be auto-formatted on save (updated user settings).',
+            )
+            break
+          case AutoFormatChoice.OnlyInWorkspace:
+            vscode.window.showInformationMessage(
+              'BAML files will now be auto-formatted on save (updated workspace settings).',
+            )
+            break
+        }
+      }),
     )
 
     activateClient(context, serverOptions, clientOptions)