diff --git a/baml_language/crates/bridge_cffi/src/baml_to_host.rs b/baml_language/crates/bridge_cffi/src/baml_to_host.rs index 1f670ac04cf..a1aedf2c82f 100644 --- a/baml_language/crates/bridge_cffi/src/baml_to_host.rs +++ b/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 { + 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, + function_name: String, + args: BexArgs, + call_ctx: FunctionCallContext, ) -> Vec { 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, + handle_key: u64, + args: BexArgs, + call_ctx: FunctionCallContext, +) -> Vec { + 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, 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; diff --git a/baml_language/crates/bridge_cffi/src/lib.rs b/baml_language/crates/bridge_cffi/src/lib.rs index 1154590432d..743525217b7 100644 --- a/baml_language/crates/bridge_cffi/src/lib.rs +++ b/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, + cancel: bex_project::CancellationToken, } -static ACTIVE_CALL_RUNTIMES: LazyLock>>> = +static ACTIVE_CALL_ROUTES: LazyLock>>> = LazyLock::new(|| Mutex::new(HashMap::new())); -fn active_call_runtimes() -> MutexGuard<'static, HashMap>> { - ACTIVE_CALL_RUNTIMES +fn active_call_routes() -> MutexGuard<'static, HashMap>> { + 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, -) -> 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 { + 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::*; diff --git a/baml_language/crates/bridge_cffi/src/lib_native.rs b/baml_language/crates/bridge_cffi/src/lib_native.rs index 952c94a0f19..edb16188484 100644 --- a/baml_language/crates/bridge_cffi/src/lib_native.rs +++ b/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 } } }) diff --git a/baml_language/sdks/rust/bridge_rust/Cargo.toml b/baml_language/sdks/rust/bridge_rust/Cargo.toml index 9982285b258..3f844fd44ed 100644 --- a/baml_language/sdks/rust/bridge_rust/Cargo.toml +++ b/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 diff --git a/baml_language/sdks/rust/bridge_rust/README.md b/baml_language/sdks/rust/bridge_rust/README.md index 8c2a683244f..4570117c3fc 100644 --- a/baml_language/sdks/rust/bridge_rust/README.md +++ b/baml_language/sdks/rust/bridge_rust/README.md @@ -8,3 +8,5 @@ BAML's typed `throws` contracts as `Result>`. 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: . + +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. diff --git a/baml_language/sdks/rust/bridge_rust/src/capi.rs b/baml_language/sdks/rust/bridge_rust/src/capi.rs index 91ce5a2570a..6ed36261cf7 100644 --- a/baml_language/sdks/rust/bridge_rust/src/capi.rs +++ b/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, call_function: Option, new_function_call: Option u64>, - /// Layout placeholder: sits between `new_function_call` and the - /// host-value entries in ABI order. Unused until cancellation lands. - cancel_function_call: Option i32>, + /// Cancels an in-flight engine call by its engine-issued id. + cancel_function_call: Option, register_host_dispatch_callback: Option, register_host_release_callback: Option, complete_host_call: Option, @@ -255,7 +259,8 @@ fn load_inner(env: &loader::LoaderEnv) -> Result { 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 { initialize_runtime_from_bytecode_with_metadata, register_callback, new_function_call, + cancel_function_call, call_function, handle_clone, handle_release, diff --git a/baml_language/sdks/rust/bridge_rust/src/completion.rs b/baml_language/sdks/rust/bridge_rust/src/completion.rs index 1601db5d7df..1ccc5bee28e 100644 --- a/baml_language/sdks/rust/bridge_rust/src/completion.rs +++ b/baml_language/sdks/rust/bridge_rust/src/completion.rs @@ -15,14 +15,16 @@ use std::{ task::{Poll, Waker}, }; -use crate::capi; +use crate::{SdkError, capi}; + +type CompletionResult = Result, SdkError>; /// One in-flight call's state. enum State { Pending, /// An async receiver parked its waker while pending. PendingWithWaker(Waker), - Ready(Vec), + Ready(CompletionResult), /// The receiver was dropped before the result arrived; the callback /// discards the payload. Abandoned, @@ -38,6 +40,8 @@ struct Slot { /// reclaimed by the callback or eagerly on drop). pub(crate) struct Receiver { dispatch_id: u32, + engine_call_id: u64, + cancel_function_call: capi::CancelFunctionCallFn, slot: Arc, } @@ -51,7 +55,15 @@ fn registry() -> &'static Mutex>> { /// Allocate a dispatch id and register a completion for it, ensuring the /// process-global callback is registered with the engine first (so a /// result can never arrive unroutable). -pub(crate) fn register(api: &'static capi::Api) -> Receiver { +pub(crate) fn register(api: &'static capi::Api, engine_call_id: u64) -> Receiver { + register_with_cancel(api, engine_call_id, api.cancel_function_call) +} + +fn register_with_cancel( + api: &'static capi::Api, + engine_call_id: u64, + cancel_function_call: capi::CancelFunctionCallFn, +) -> Receiver { static CALLBACK_REGISTERED: OnceLock<()> = OnceLock::new(); // Dispatch ids only correlate callback deliveries with waiting // receivers; wrap-around is harmless as long as ~4 billion calls are @@ -77,7 +89,12 @@ pub(crate) fn register(api: &'static capi::Api) -> Receiver { .lock() .expect("completion registry poisoned") .insert(dispatch_id, Arc::clone(&slot)); - Receiver { dispatch_id, slot } + Receiver { + dispatch_id, + engine_call_id, + cancel_function_call, + slot, + } } impl Receiver { @@ -86,14 +103,14 @@ impl Receiver { } /// Block until the result envelope arrives. - pub(crate) fn wait_blocking(self) -> Vec { + pub(crate) fn wait_blocking(self) -> CompletionResult { let mut state = self.slot.state.lock().expect("completion slot poisoned"); loop { match std::mem::replace(&mut *state, State::Pending) { - State::Ready(bytes) => { + State::Ready(result) => { *state = State::Abandoned; drop(state); - return bytes; + return result; } other => { *state = other; @@ -108,20 +125,20 @@ impl Receiver { } /// Await the result envelope. The future is executor-agnostic. - pub(crate) async fn wait(self) -> Vec { + pub(crate) async fn wait(self) -> CompletionResult { struct WaitFuture(Receiver); impl Future for WaitFuture { - type Output = Vec; + type Output = CompletionResult; fn poll( self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, - ) -> Poll> { + ) -> Poll { let mut state = self.0.slot.state.lock().expect("completion slot poisoned"); match std::mem::replace(&mut *state, State::Pending) { - State::Ready(bytes) => { + State::Ready(result) => { *state = State::Abandoned; - Poll::Ready(bytes) + Poll::Ready(result) } State::Pending | State::PendingWithWaker(_) => { *state = State::PendingWithWaker(cx.waker().clone()); @@ -141,12 +158,30 @@ impl Drop for Receiver { // an entry left in the registry means the result never arrived (or // the receiver is being dropped unconsumed) — reclaim it so the // table cannot grow without bound. - registry() + let was_registered = registry() .lock() .expect("completion registry poisoned") - .remove(&self.dispatch_id); + .remove(&self.dispatch_id) + .is_some(); let mut state = self.slot.state.lock().expect("completion slot poisoned"); + let was_pending = matches!(*state, State::Pending | State::PendingWithWaker(_)); *state = State::Abandoned; + drop(state); + + // A registered, pending receiver represents a caller that stopped + // observing the call (for example, `tokio::time::timeout` dropped its + // future). Propagate that cancellation to the engine before the + // caller's own concurrency permit can be reused. If the callback won + // the race and removed the registry entry, the engine call is already + // complete and must not be cancelled. + if was_registered && was_pending { + // SAFETY: the function pointer came from the validated process- + // lifetime C API table and accepts this engine-issued call id. + #[expect(unsafe_code)] + unsafe { + (self.cancel_function_call)(self.engine_call_id); + } + } } } @@ -155,13 +190,25 @@ impl Drop for Receiver { /// dispatch ids are discarded. Must never unwind into the engine. extern "C" fn trampoline(call_id: u32, content: *const c_char, length: usize) { let caught = std::panic::catch_unwind(|| { - // SAFETY: the engine guarantees `content` is valid for `length` - // bytes for the synchronous duration of this call. - let bytes = if content.is_null() { - Vec::new() + let result = if length > crate::runtime::MAX_RESULT_BYTES { + Err(SdkError::new(format!( + "BAML result exceeded the {} MiB bridge limit (received {length} bytes)", + crate::runtime::MAX_RESULT_BYTES / (1024 * 1024), + ))) + } else if content.is_null() && length != 0 { + Err(SdkError::new( + "engine returned a null BAML result pointer with a nonzero length", + )) } else { - #[expect(unsafe_code)] - unsafe { std::slice::from_raw_parts(content.cast::(), length) }.to_vec() + // SAFETY: the engine guarantees `content` is valid for `length` + // bytes for the synchronous duration of this call. + let bytes = if length == 0 { + Vec::new() + } else { + #[expect(unsafe_code)] + unsafe { std::slice::from_raw_parts(content.cast::(), length) }.to_vec() + }; + Ok(bytes) }; let slot = registry() .lock() @@ -169,7 +216,7 @@ extern "C" fn trampoline(call_id: u32, content: *const c_char, length: usize) { .remove(&call_id); if let Some(slot) = slot { let mut state = slot.state.lock().expect("completion slot poisoned"); - let previous = std::mem::replace(&mut *state, State::Ready(bytes)); + let previous = std::mem::replace(&mut *state, State::Ready(result)); drop(state); match previous { State::PendingWithWaker(waker) => waker.wake(), @@ -195,11 +242,16 @@ extern "C" fn trampoline(call_id: u32, content: *const c_char, length: usize) { mod tests { use super::*; - /// Register against the real loaded engine's table; these tests - /// fulfill through the trampoline directly and never call it. + /// Register against the real loaded engine's table, but stub cancellation; + /// these tests fulfill through the trampoline directly and never start an + /// engine call. fn register() -> Receiver { crate::test_support::locate_dev_engine(); - super::register(capi::api().expect("engine library loads")) + register_with_cancel( + capi::api().expect("engine library loads"), + 1, + record_cancellation, + ) } // Fulfill directly through the trampoline, as the engine would. @@ -216,14 +268,14 @@ mod tests { // on the ordering either way. std::thread::yield_now(); fulfill(id, b"hello"); - assert_eq!(handle.join().unwrap(), b"hello"); + assert_eq!(handle.join().unwrap().unwrap(), b"hello"); } #[test] fn fulfill_before_wait_is_immediate() { let receiver = register(); fulfill(receiver.dispatch_id(), b"early"); - assert_eq!(receiver.wait_blocking(), b"early"); + assert_eq!(receiver.wait_blocking().unwrap(), b"early"); } #[test] @@ -243,7 +295,58 @@ mod tests { let handle = std::thread::spawn(move || minimal_block_on(receiver.wait())); std::thread::yield_now(); fulfill(id, b"async"); - assert_eq!(handle.join().unwrap(), b"async"); + assert_eq!(handle.join().unwrap().unwrap(), b"async"); + } + + fn fake_cancellations() -> &'static Mutex> { + static CALLS: OnceLock>> = OnceLock::new(); + CALLS.get_or_init(|| Mutex::new(Vec::new())) + } + + #[expect(unsafe_code, reason = "matches the engine cancellation ABI")] + unsafe extern "C" fn record_cancellation(call_id: u64) -> i32 { + fake_cancellations().lock().unwrap().push(call_id); + 0 + } + + fn fake_cancel_receiver(engine_call_id: u64) -> Receiver { + crate::test_support::locate_dev_engine(); + register_with_cancel( + capi::api().expect("engine library loads"), + engine_call_id, + record_cancellation, + ) + } + + #[test] + fn dropping_pending_receiver_cancels_its_engine_call() { + const ENGINE_CALL_ID: u64 = 0xCA11_CE11; + let receiver = fake_cancel_receiver(ENGINE_CALL_ID); + drop(receiver); + let calls = fake_cancellations().lock().unwrap(); + assert_eq!(calls.iter().filter(|&&id| id == ENGINE_CALL_ID).count(), 1); + } + + #[test] + fn dropping_fulfilled_receiver_does_not_cancel() { + const ENGINE_CALL_ID: u64 = 0xC0DE_0001; + let receiver = fake_cancel_receiver(ENGINE_CALL_ID); + fulfill(receiver.dispatch_id(), b"done"); + drop(receiver); + let calls = fake_cancellations().lock().unwrap(); + assert!(!calls.contains(&ENGINE_CALL_ID)); + } + + #[test] + fn oversized_result_is_rejected_without_reading_its_payload() { + let receiver = register(); + trampoline( + receiver.dispatch_id(), + std::ptr::dangling(), + crate::runtime::MAX_RESULT_BYTES + 1, + ); + let error = receiver.wait_blocking().unwrap_err(); + assert!(error.to_string().contains("32 MiB bridge limit")); } /// Minimal single-future `block_on` so the test needs no async runtime. diff --git a/baml_language/sdks/rust/bridge_rust/src/runtime.rs b/baml_language/sdks/rust/bridge_rust/src/runtime.rs index 74cc3cbfb3a..73946fdb0ae 100644 --- a/baml_language/sdks/rust/bridge_rust/src/runtime.rs +++ b/baml_language/sdks/rust/bridge_rust/src/runtime.rs @@ -13,6 +13,13 @@ use prost::Message as _; use crate::{BamlValue, Error, SdkError, capi, completion, decode, wire}; +/// Maximum encoded result envelope accepted from the engine. +/// +/// This preserves the 32 MiB output ceiling of the former CLI subprocess +/// boundary and prevents the callback boundary from copying an unbounded +/// payload into the caller. +pub const MAX_RESULT_BYTES: usize = 32 * 1024 * 1024; + /// Initialize (or replace) the process-global runtime from the /// borsh-encoded bytecode a generated SDK embeds. /// @@ -91,10 +98,9 @@ pub fn invoke_sync( } let receiver = dispatch(fqn, kwargs, type_args).map_err(Error::Sdk)?; // Blocks until the engine delivers the result envelope via the callback. - // There is no timeout: the engine is contracted to complete every call - // (success, thrown error, or panic). A caller-facing timeout/cancellation - // path lands with the cancellation feature (`cancel_function_call`). - let bytes = receiver.wait_blocking(); + // Synchronous calls wait for engine completion; async calls propagate + // future cancellation automatically. + let bytes = receiver.wait_blocking().map_err(Error::Sdk)?; decode::decode_result(&bytes) } @@ -109,7 +115,7 @@ pub async fn invoke( type_args: Vec, ) -> Result> { let receiver = dispatch(fqn, kwargs, type_args).map_err(Error::Sdk)?; - let bytes = receiver.wait().await; + let bytes = receiver.wait().await.map_err(Error::Sdk)?; decode::decode_result(&bytes) } @@ -121,7 +127,8 @@ pub fn invoke_handle_sync( return Err(Error::CalledSyncFromAsync); } let receiver = dispatch_handle(handle_key, kwargs).map_err(Error::Sdk)?; - decode::decode_result(&receiver.wait_blocking()) + let bytes = receiver.wait_blocking().map_err(Error::Sdk)?; + decode::decode_result(&bytes) } pub async fn invoke_handle( @@ -129,7 +136,8 @@ pub async fn invoke_handle( kwargs: Vec, ) -> Result> { let receiver = dispatch_handle(handle_key, kwargs).map_err(Error::Sdk)?; - decode::decode_result(&receiver.wait().await) + let bytes = receiver.wait().await.map_err(Error::Sdk)?; + decode::decode_result(&bytes) } /// Encode the call and fire it through the C ABI. The registered @@ -141,7 +149,6 @@ fn dispatch( type_args: Vec, ) -> Result { let api = capi::api()?; - let receiver = completion::register(api); // Host-callable dispatch must be installed before the engine can hold // a callable handle; every handle rides a call that passes through // here first. @@ -149,6 +156,7 @@ fn dispatch( // SAFETY: takes no arguments; allocates an id inside the engine. #[expect(unsafe_code)] let call_id = unsafe { (api.new_function_call)() }; + let receiver = completion::register(api, call_id); let args = wire::CallFunctionArgs { kwargs, call_id, @@ -174,12 +182,12 @@ fn dispatch_handle( return Err(SdkError::new("cannot invoke a zero BAML function handle")); } let api = capi::api()?; - let receiver = completion::register(api); crate::host_value::ensure_callbacks_registered(api); // SAFETY: this ABI function takes no arguments and returns a fresh engine // call id; the loaded API table was layout-checked during initialization. #[expect(unsafe_code)] let call_id = unsafe { (api.new_function_call)() }; + let receiver = completion::register(api, call_id); let args = wire::CallFunctionArgs { kwargs, call_id, diff --git a/baml_language/sdks/rust/bridge_rust/tests/live_engine.rs b/baml_language/sdks/rust/bridge_rust/tests/live_engine.rs index e8c1530da02..f5436b37b75 100644 --- a/baml_language/sdks/rust/bridge_rust/tests/live_engine.rs +++ b/baml_language/sdks/rust/bridge_rust/tests/live_engine.rs @@ -2,7 +2,15 @@ //! inline BAML source → compile → invoke → decode, exercising the same //! `bridge_cffi` / `bridge_ctypes` machinery a generated SDK composes. -use std::{collections::HashMap, convert::Infallible, sync::OnceLock}; +use std::{ + collections::HashMap, + convert::Infallible, + sync::{ + Arc, OnceLock, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; mod common; @@ -21,6 +29,10 @@ function no_op() -> void {} function opt_probe(a: int, o: int? = 5) -> int?[] { [a, o] } function gid(x: T) -> T { x } function gname() -> string { type.of().to_string() } +function slow_callback(callback: () -> void, ms: int) -> void { + baml.sys.sleep(baml.time.Duration.from_milliseconds(ms)); + callback() +} "#; /// Initialize the process-global runtime once for every test in this @@ -163,6 +175,40 @@ async fn async_invoke_round_trips() { assert_eq!(result, 42); } +#[tokio::test] +async fn dropping_timed_out_invoke_cancels_the_engine_call() { + ensure_runtime(); + let called = Arc::new(AtomicBool::new(false)); + let called_from_callback = Arc::clone(&called); + let callback = baml_bridge::host_value::callable_handle( + move || called_from_callback.store(true, Ordering::SeqCst), + &[], + ); + + let result = tokio::time::timeout( + Duration::from_millis(50), + runtime::invoke::<(), Infallible>( + "user.slow_callback", + encode::kwargs(vec![ + ("callback", Some(callback)), + ("ms", Some(500i64.to_baml())), + ]), + vec![], + ), + ) + .await; + assert!( + result.is_err(), + "the call should exceed the caller deadline" + ); + + // If dropping the timed-out future merely detached the engine call, it + // would wake after 500 ms and invoke the host callback. Cancellation must + // stop it before that observable side effect. + tokio::time::sleep(Duration::from_millis(600)).await; + assert!(!called.load(Ordering::SeqCst)); +} + #[tokio::test] async fn sync_invoke_inside_async_runtime_is_refused() { ensure_runtime();