fix(rust): cancel dropped bridge calls

baml_language/crates/bridge_cffi/src/baml_to_host.rs

@@ -324,9 +324,23 @@ pub async fn call_and_encode(
     function_name: String,
     args: BexArgs,
     call_ctx: FunctionCallContext,
+) -> Vec<u8> {
+    let _route =
+        match crate::register_active_call_route(call_ctx.host_call_id.0, call_ctx.cancel.clone()) {
+            Ok(route) => route,
+            Err(error) => return error_to_outbound(error),
+        };
+    call_and_encode_registered(runtime, function_name, args, call_ctx).await
+}
+
+/// Named-call encoding for callers that already own an active-call route.
+pub(crate) async fn call_and_encode_registered(
+    runtime: Arc<dyn Bex>,
+    function_name: String,
+    args: BexArgs,
+    call_ctx: FunctionCallContext,
 ) -> Vec<u8> {
     let options = CffiHandleTableOptions::for_in_process();
-    let _route = crate::register_active_call_runtime(call_ctx.host_call_id.0, &runtime);
 
     let caught = AssertUnwindSafe(runtime.call_function(&function_name, args, call_ctx))
         .catch_unwind()
@@ -372,6 +386,21 @@ fn partition_callable_args(
 /// Invoke an engine-owned callable referenced by an ordinary handle-table key
 /// and encode the result through the same envelope path as a named call.
 pub async fn call_handle_and_encode(
+    runtime: Arc<dyn Bex>,
+    handle_key: u64,
+    args: BexArgs,
+    call_ctx: FunctionCallContext,
+) -> Vec<u8> {
+    let _route =
+        match crate::register_active_call_route(call_ctx.host_call_id.0, call_ctx.cancel.clone()) {
+            Ok(route) => route,
+            Err(error) => return error_to_outbound(error),
+        };
+    call_handle_and_encode_registered(runtime, handle_key, args, call_ctx).await
+}
+
+/// Handle-call encoding for callers that already own an active-call route.
+pub(crate) async fn call_handle_and_encode_registered(
     runtime: Arc<dyn Bex>,
     handle_key: u64,
     BexArgs { required, optional }: BexArgs,
@@ -414,7 +443,6 @@ pub async fn call_handle_and_encode(
     };
 
     let options = CffiHandleTableOptions::for_in_process();
-    let _route = crate::register_active_call_runtime(call_ctx.host_call_id.0, &runtime);
     let caught = AssertUnwindSafe(runtime.call_callable(handle, args, call_ctx))
         .catch_unwind()
         .await;

baml_language/crates/bridge_cffi/src/lib.rs

@@ -16,7 +16,7 @@ mod platform;
 
 use std::{
     collections::HashMap,
-    sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError, Weak},
+    sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError},
 };
 
 use bex_project::Bex;
@@ -91,14 +91,14 @@ fn source_vfs_root() -> vfs::VfsPath {
 }
 
 struct ActiveCallRoute {
-    runtime: Weak<dyn Bex>,
+    cancel: bex_project::CancellationToken,
 }
 
-static ACTIVE_CALL_RUNTIMES: LazyLock<Mutex<HashMap<u64, Arc<ActiveCallRoute>>>> =
+static ACTIVE_CALL_ROUTES: LazyLock<Mutex<HashMap<u64, Arc<ActiveCallRoute>>>> =
     LazyLock::new(|| Mutex::new(HashMap::new()));
 
-fn active_call_runtimes() -> MutexGuard<'static, HashMap<u64, Arc<ActiveCallRoute>>> {
-    ACTIVE_CALL_RUNTIMES
+fn active_call_routes() -> MutexGuard<'static, HashMap<u64, Arc<ActiveCallRoute>>> {
+    ACTIVE_CALL_ROUTES
         .lock()
         .unwrap_or_else(PoisonError::into_inner)
 }
@@ -110,7 +110,7 @@ pub(crate) struct ActiveCallRouteGuard {
 
 impl Drop for ActiveCallRouteGuard {
     fn drop(&mut self) {
-        let mut routes = active_call_runtimes();
+        let mut routes = active_call_routes();
         if routes
             .get(&self.call_id)
             .is_some_and(|current| Arc::ptr_eq(current, &self.route))
@@ -120,15 +120,21 @@ impl Drop for ActiveCallRouteGuard {
     }
 }
 
-pub(crate) fn register_active_call_runtime(
+pub(crate) fn register_active_call_route(
     call_id: u64,
-    runtime: &Arc<dyn Bex>,
-) -> ActiveCallRouteGuard {
-    let route = Arc::new(ActiveCallRoute {
-        runtime: Arc::downgrade(runtime),
-    });
-    active_call_runtimes().insert(call_id, Arc::clone(&route));
-    ActiveCallRouteGuard { call_id, route }
+    cancel: bex_project::CancellationToken,
+) -> Result<ActiveCallRouteGuard, BridgeError> {
+    if call_id == 0 {
+        return Err(BridgeError::InvalidCallId);
+    }
+    let route = Arc::new(ActiveCallRoute { cancel });
+    let mut routes = active_call_routes();
+    if routes.contains_key(&call_id) {
+        return Err(BridgeError::DuplicateCallId(call_id));
+    }
+    routes.insert(call_id, Arc::clone(&route));
+    drop(routes);
+    Ok(ActiveCallRouteGuard { call_id, route })
 }
 
 pub mod baml_to_host;
@@ -400,17 +406,22 @@ pub fn function_call_context_builder(
 
 /// Cancel an in-flight function call by ID.
 ///
-/// Returns true on success, false if the runtime is not initialized.
+/// Returns true on success, false if no route or initialized runtime exists.
 pub fn cancel_function_call_by_id(id: u64) -> bool {
     if id == 0 {
         return false;
     }
-    let originating_runtime = active_call_runtimes()
+    // Routed calls own their token until result delivery finishes. Cancelling
+    // it directly is safe both before engine registration and after engine
+    // completion; the latter must not create an unconsumed pre-cancel entry.
+    if let Some(cancel) = active_call_routes()
         .get(&id)
-        .and_then(|route| route.runtime.upgrade());
-    originating_runtime
-        .map(Ok)
-        .unwrap_or_else(get_runtime)
+        .map(|route| route.cancel.clone())
+    {
+        cancel.cancel();
+        return true;
+    }
+    get_runtime()
         .and_then(|runtime| {
             runtime
                 .cancel_function_call(bex_project::CallId(id))
@@ -433,6 +444,59 @@ pub extern "C" fn cancel_function_call(id: u64) -> i32 {
     if cancel_function_call_by_id(id) { 0 } else { 1 }
 }
 
+#[cfg(test)]
+mod cancellation_route_tests {
+    use super::*;
+
+    #[test]
+    fn registered_route_cancels_its_token_without_reserving_an_engine_call() {
+        let call_id = new_function_call_id();
+        let cancel = bex_project::CancellationToken::new();
+        let route = register_active_call_route(call_id, cancel.clone()).unwrap();
+
+        assert!(cancel_function_call_by_id(call_id));
+        assert!(cancel.is_cancelled());
+        assert!(active_call_routes().contains_key(&call_id));
+
+        drop(route);
+        assert!(!active_call_routes().contains_key(&call_id));
+    }
+
+    #[test]
+    fn duplicate_route_is_rejected_without_replacing_the_original() {
+        let call_id = new_function_call_id();
+        let first_cancel = bex_project::CancellationToken::new();
+        let second_cancel = bex_project::CancellationToken::new();
+        let first_route = register_active_call_route(call_id, first_cancel.clone()).unwrap();
+
+        let error = match register_active_call_route(call_id, second_cancel.clone()) {
+            Ok(_) => panic!("duplicate route unexpectedly replaced the original"),
+            Err(error) => error,
+        };
+        assert!(matches!(error, BridgeError::DuplicateCallId(id) if id == call_id));
+
+        assert!(cancel_function_call_by_id(call_id));
+        assert!(first_cancel.is_cancelled());
+        assert!(!second_cancel.is_cancelled());
+        drop(first_route);
+    }
+
+    #[test]
+    fn zero_route_is_rejected_without_registration() {
+        let cancel = bex_project::CancellationToken::new();
+
+        let error = match register_active_call_route(0, cancel.clone()) {
+            Ok(_) => panic!("zero route unexpectedly registered"),
+            Err(error) => error,
+        };
+
+        assert!(matches!(error, BridgeError::InvalidCallId));
+        assert!(!active_call_routes().contains_key(&0));
+        assert!(!cancel.is_cancelled());
+        assert!(!cancel_function_call_by_id(0));
+    }
+}
+
 #[cfg(test)]
 mod generated_metadata_tests {
     use super::*;

baml_language/crates/bridge_cffi/src/lib_native.rs

@@ -14,8 +14,8 @@ use sys_native::SysOpsExt;
 use tokio::runtime::Runtime;
 
 use crate::{
-    BridgeError, baml_to_host, call_and_encode, call_handle_and_encode, error_to_outbound,
-    function_call_context_builder,
+    BridgeError, baml_to_host, error_to_outbound, function_call_context_builder,
+    register_active_call_route,
 };
 
 #[path = "api.rs"]
@@ -159,19 +159,38 @@ fn call_function_inner(encoded_args: *const u8, length: usize, id: u32) -> Resul
     }
     let type_args = bridge_ctypes::proto_ty_args_to_named(&args.type_args)?;
     let kwargs = kwargs_to_bex_values(args.kwargs, &HANDLE_TABLE)?;
+    let cancel = bex_project::CancellationToken::new();
     let call_ctx = function_call_context_builder(call_id)
+        .with_cancel_token(cancel.clone())
         .with_type_args(type_args.type_args)
-        .with_type_defs(type_args.type_defs);
+        .with_type_defs(type_args.type_defs)
+        .build();
+    // Install the cancellation route before returning control to the caller.
+    // The guard stays alive through synchronous callback delivery, closing
+    // both the pre-start and completed-before-delivery races.
+    let route = register_active_call_route(call_id.0, cancel)?;
 
     get_tokio_runtime()?.spawn(async move {
+        let _route = route;
         let encoded = AssertUnwindSafe(async move {
             match target {
                 CallTarget::FunctionName(function_name) => {
-                    call_and_encode(runtime, function_name, kwargs.into(), call_ctx.build()).await
+                    baml_to_host::call_and_encode_registered(
+                        runtime,
+                        function_name,
+                        kwargs.into(),
+                        call_ctx,
+                    )
+                    .await
                 }
                 CallTarget::FunctionHandle(handle_key) => {
-                    call_handle_and_encode(runtime, handle_key, kwargs.into(), call_ctx.build())
-                        .await
+                    baml_to_host::call_handle_and_encode_registered(
+                        runtime,
+                        handle_key,
+                        kwargs.into(),
+                        call_ctx,
+                    )
+                    .await
                 }
             }
         })

baml_language/sdks/rust/bridge_rust/Cargo.toml

@@ -42,7 +42,7 @@ tokio = { workspace = true, features = [ "rt", "rt-multi-thread" ] }
 ureq = { workspace = true }
 
 [dev-dependencies]
-tokio = { workspace = true, features = [ "rt", "rt-multi-thread", "macros" ] }
+tokio = { workspace = true, features = [ "rt", "rt-multi-thread", "macros", "time" ] }
 
 [lints]
 workspace = true

baml_language/sdks/rust/bridge_rust/README.md

@@ -8,3 +8,5 @@ BAML's typed `throws` contracts as `Result<T, baml_bridge::Error<E>>`.
 You normally don't add this crate by hand — the generated `baml_sdk` crate
 pins the matching version. See the BAML documentation for getting started:
 <https://docs.boundaryml.com>.
+
+Async calls are cancellation-safe: dropping a generated `_async` future (for example, when `tokio::time::timeout` expires) cancels the corresponding engine call instead of leaving it running detached. Completed result envelopes are limited to 32 MiB at the bridge boundary.

baml_language/sdks/rust/bridge_rust/src/capi.rs

@@ -18,6 +18,10 @@ use crate::{
 /// the synchronous duration of the call — implementations must copy.
 pub(crate) type CallbackFn = extern "C" fn(call_id: u32, content: *const c_char, length: usize);
 
+/// Request cancellation of one engine call. Zero means the cancellation was
+/// accepted; nonzero means the call was unknown or already complete.
+pub(crate) type CancelFunctionCallFn = unsafe extern "C" fn(u64) -> i32;
+
 /// The engine's BAML→host dispatch callback: BAML invoked a host-owned
 /// callable. `args` (a protobuf `BamlToHostCall`) is borrowed only for the
 /// synchronous duration of the call — implementations must copy, return
@@ -50,6 +54,7 @@ pub(crate) struct Api {
         unsafe extern "C" fn(*const u8, usize, *const c_char) -> Buffer,
     pub(crate) register_callback: unsafe extern "C" fn(CallbackFn),
     pub(crate) new_function_call: unsafe extern "C" fn() -> u64,
+    pub(crate) cancel_function_call: CancelFunctionCallFn,
     pub(crate) call_function: unsafe extern "C" fn(*const u8, usize, u32),
     pub(crate) handle_clone: unsafe extern "C" fn(u64, *mut u64) -> u32,
     pub(crate) handle_release: unsafe extern "C" fn(u64) -> u32,
@@ -144,9 +149,8 @@ struct BamlApiV1 {
     register_callback: Option<unsafe extern "C" fn(CallbackFn)>,
     call_function: Option<unsafe extern "C" fn(*const u8, usize, u32)>,
     new_function_call: Option<unsafe extern "C" fn() -> u64>,
-    /// Layout placeholder: sits between `new_function_call` and the
-    /// host-value entries in ABI order. Unused until cancellation lands.
-    cancel_function_call: Option<unsafe extern "C" fn(u64) -> i32>,
+    /// Cancels an in-flight engine call by its engine-issued id.
+    cancel_function_call: Option<CancelFunctionCallFn>,
     register_host_dispatch_callback: Option<unsafe extern "C" fn(HostDispatchFn)>,
     register_host_release_callback: Option<unsafe extern "C" fn(HostReleaseFn)>,
     complete_host_call: Option<unsafe extern "C" fn(u32, i32, *const c_char, usize)>,
@@ -255,7 +259,8 @@ fn load_inner(env: &loader::LoaderEnv) -> Result<Api, LoaderError> {
     let register_callback = required_slot(table.register_callback, "register_callback", &path)?;
     let call_function = required_slot(table.call_function, "call_function", &path)?;
     let new_function_call = required_slot(table.new_function_call, "new_function_call", &path)?;
-    required_slot(table.cancel_function_call, "cancel_function_call", &path)?;
+    let cancel_function_call =
+        required_slot(table.cancel_function_call, "cancel_function_call", &path)?;
     let register_host_dispatch_callback = required_slot(
         table.register_host_dispatch_callback,
         "register_host_dispatch_callback",
@@ -291,6 +296,7 @@ fn load_inner(env: &loader::LoaderEnv) -> Result<Api, LoaderError> {
         initialize_runtime_from_bytecode_with_metadata,
         register_callback,
         new_function_call,
+        cancel_function_call,
         call_function,
         handle_clone,
         handle_release,