feat(serve): Add graceful Ctrl+C handling with clean exit message

engine/baml-runtime/Cargo.toml

@@ -65,6 +65,7 @@ web-time.workspace = true
 static_assertions.workspace = true
 mime_guess = "=2.0.5"
 mime = "0.3.17"
+libc = "0.2"
 
 # For tracing
 envy = "0.4.2"
@@ -104,7 +105,7 @@ colored = { version = "2.1.0", default-features = false, features = [
 futures-timer = { version = "3.0.3", features = ["wasm-bindgen"] }
 js-sys = "0.3.69"
 reqwest = { version = "0.12.5", features = ["stream", "json"] }
-# 
+#
 send_wrapper = { version = "0.6.0", features = ["futures"] }
 serde-wasm-bindgen = "0.6.5"
 uuid = { version = "1.8.0", features = ["v4", "serde", "js"] }
@@ -137,6 +138,7 @@ tokio = { version = "1", features = ["full"] }
 reqwest.workspace = true
 walkdir = "2.5.0"
 which = "6.0.3"
+libc = "0.2"
 
 
 

engine/baml-runtime/src/cli/serve/mod.rs

@@ -6,6 +6,10 @@ use error::BamlError;
 use indexmap::IndexMap;
 use internal_baml_codegen::GeneratorArgs;
 use json_response::Json;
+use std::io::Write;
+
+#[cfg(all(unix, not(target_arch = "wasm32")))]
+use libc;
 
 use anyhow::{Context, Result};
 use arg_validation::BamlServeValidate;
@@ -66,7 +70,7 @@ impl ServeArgs {
         if !self.preview {
             log::warn!(
                 r#"BAML-over-HTTP API is a preview feature.
-                
+
 Please run with --preview, like so:
 
     {} serve --preview
@@ -313,21 +317,7 @@ impl Server {
             tcp_listener,
             app.layer(axum::middleware::from_fn(Server::auth_middleware)),
         );
-        // TODO: we do not handle this ourselves, because tokio's default
-        // handling is pretty good on unix.
-        //
-        // Not totally sure if this WAI on
-        // windows, but there are some potential pitfalls that we can run into
-        // if we try to handle this ourselves, see
-        // https://docs.rs/tokio/latest/tokio/signal/fn.ctrl_c.html#caveats
-        //
-        // Namely- we need to ensure resilient delivery of Ctrl-C to everything, and I
-        // suspect we need to do a bit of work to ensure that we handle that bookkeeping
-        // correctly. Shutting down the BAML runtime, tokio runtime, _and_ axum is not
-        // super straightforward, because I don't know how much is handled for us
-        // out of the box.
-        //
-        // .with_graceful_shutdown(signal::ctrl_c());
+
         log::info!(
             r#"BAML-over-HTTP listening on port {}, serving from {}
 
@@ -340,7 +330,29 @@ Tip: test that the server is up using `curl http://localhost:{}/_debug/ping`
             self.port,
         );
 
-        service.await?;
+        // Add graceful shutdown with Ctrl+C handling, and a clean message for exit to the client
+        let shutdown = async {
+            match tokio::signal::ctrl_c().await {
+                Ok(()) => {
+                    // Print message and ensure it's flushed
+                    println!("\nExiting BAML...");
+                    std::io::stdout().flush().unwrap_or_default();
+                    // Small delay to ensure message is visible
+                    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+                }
+                Err(err) => {
+                    log::error!("Error handling Ctrl+C: {}", err);
+                }
+            }
+        };
+
+        // Set ctrl-c handler to ignore so Python doesn't intercept it, otherwise python doesnt like it
+        #[cfg(all(unix, not(target_arch = "wasm32")))]
+        unsafe {
+            libc::signal(libc::SIGINT, libc::SIG_IGN);
+        }
+
+        service.with_graceful_shutdown(shutdown).await?;
 
         Ok(())
     }