bridge(feat): track created handles and enable batch cleanup#3195
bridge(feat): track created handles and enable batch cleanup#3195miguelcsx wants to merge 1 commit intoBoundaryML:canaryfrom
Conversation
- track created handles in outbound values so runtimes can release unused ones - add batch release helpers across cffi/python/wasm - add media inlining guard and handle-type validation to reduce unsafe/mismatched usage
|
@miguelcsx is attempting to deploy a commit to the Boundary Team on Vercel. A member of the Team first needs to authorize it. |
📝 WalkthroughWalkthroughThis pull request introduces bulk handle release functionality across multiple FFI bridges (CFFI, Python, WASM) and adds handle type validation during decoding. It refactors the outbound value encoding pipeline to track created handles, adds a configurable media size limit, and extends the HandleTable with a Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
baml_language/crates/bridge_ctypes/src/value_encode.rs (1)
153-159:⚠️ Potential issue | 🟠 MajorPromptAst media bypasses
max_inline_media_bytesand can still inline oversized payloads.When
serialize_prompt_astis enabled, media insidePromptAstSimple::Mediais always inlined viabex_media_to_proto_media, so large base64 blobs can cross the boundary despite the new guard.🛠️ Suggested fix direction
- BexExternalValue::Adt(BexExternalAdt::PromptAst(prompt_ast)) - if options.serialize_prompt_ast => + BexExternalValue::Adt(BexExternalAdt::PromptAst(prompt_ast)) + if options.serialize_prompt_ast + && prompt_ast_within_inline_media_limit(prompt_ast, options) => { Some(BamlValueVariant::PromptAstValue( bex_prompt_ast_to_proto_prompt_ast(prompt_ast), )) }+fn prompt_ast_within_inline_media_limit( + prompt_ast: &bex_project::PromptAst, + options: &HandleTableOptions, +) -> bool { + match prompt_ast { + bex_project::PromptAst::Simple(simple) => { + prompt_ast_simple_within_inline_media_limit(simple, options) + } + bex_project::PromptAst::Message { content, .. } => { + prompt_ast_simple_within_inline_media_limit(content, options) + } + bex_project::PromptAst::Vec(items) => items + .iter() + .all(|item| prompt_ast_within_inline_media_limit(item.as_ref(), options)), + } +} + +fn prompt_ast_simple_within_inline_media_limit( + simple: &bex_project::PromptAstSimple, + options: &HandleTableOptions, +) -> bool { + match simple { + bex_project::PromptAstSimple::String(_) => true, + bex_project::PromptAstSimple::Media(media) => should_inline_media(media, options), + bex_project::PromptAstSimple::Multiple(items) => items + .iter() + .all(|item| prompt_ast_simple_within_inline_media_limit(item.as_ref(), options)), + } +}Also applies to: 278-304
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (13)
baml_language/crates/bridge_cffi/src/ffi/handle.rsbaml_language/crates/bridge_cffi/src/lib.rsbaml_language/crates/bridge_ctypes/src/error.rsbaml_language/crates/bridge_ctypes/src/handle_table.rsbaml_language/crates/bridge_ctypes/src/value_decode.rsbaml_language/crates/bridge_ctypes/src/value_encode.rsbaml_language/crates/bridge_ctypes/types/baml/cffi/v1/baml_outbound.protobaml_language/crates/bridge_python/python_src/baml/cffi/v1/baml_inbound_pb2.pybaml_language/crates/bridge_python/python_src/baml/cffi/v1/baml_outbound_pb2.pybaml_language/crates/bridge_python/python_src/baml_py/proto.pybaml_language/crates/bridge_python/src/handle.rsbaml_language/crates/bridge_wasm/src/handle.rsbaml_language/crates/bridge_wasm/src/lib.rs
| /// Release multiple handles. | ||
| pub fn release_many<I>(&self, keys: I) | ||
| where | ||
| I: IntoIterator<Item = u64>, | ||
| { | ||
| let mut entries = self | ||
| .entries | ||
| .write() | ||
| .unwrap_or_else(std::sync::PoisonError::into_inner); | ||
| for key in keys { | ||
| entries.remove(&key); | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Add a dedicated unit test for release_many behavior.
The method is straightforward, but this is new cleanup-critical API surface and should have explicit coverage (present keys removed, missing keys ignored, duplicates tolerated).
🧪 Suggested unit test
#[test]
fn double_release_returns_false() {
@@
}
+#[test]
+fn release_many_releases_present_and_ignores_missing() {
+ let table = HandleTable::new();
+ let key1 = table.insert(make_function_ref());
+ let key2 = table.insert(make_function_ref());
+
+ table.release_many([key1, 9_999_999, key2, key1]);
+
+ assert!(table.resolve(key1).is_none());
+ assert!(table.resolve(key2).is_none());
+}As per coding guidelines: "Prefer writing Rust unit tests over integration tests where possible."
| #[test] | ||
| fn handle_type_mismatch_errors() { | ||
| let table = HandleTable::new(); | ||
| let key = table.insert(HandleTableValue::FunctionRef { global_index: 3 }); | ||
| let inbound = inbound_handle(key, BamlHandleType::AdtMediaImage); | ||
| let err = inbound_to_external(inbound, &table).unwrap_err(); | ||
| match err { | ||
| CtypesError::InvalidHandleType { key: err_key, .. } => { | ||
| assert_eq!(err_key, key); | ||
| } | ||
| _ => panic!("expected InvalidHandleType"), | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Add a unit test for unrecognized numeric handle types.
The permissive Err(_) => Ok(()) branch in validate_handle_type is currently untested. Add a case using a raw non-enum integer (e.g., i32::MAX) so this compatibility contract doesn’t regress.
💡 Suggested test addition
#[test]
fn handle_type_unknown_is_allowed() {
@@
}
+#[test]
+fn handle_type_unrecognized_numeric_is_allowed() {
+ let table = HandleTable::new();
+ let key = table.insert(HandleTableValue::FunctionRef { global_index: 4 });
+ let inbound = InboundValue {
+ value: Some(InboundValueVariant::Handle(BamlHandle {
+ key,
+ handle_type: i32::MAX,
+ })),
+ };
+ let out = inbound_to_external(inbound, &table).unwrap();
+ assert!(matches!(
+ out,
+ BexExternalValue::FunctionRef { global_index: 4 }
+ ));
+}As per coding guidelines: "Prefer writing Rust unit tests over integration tests where possible."
| #[wasm_bindgen(js_name = "releaseHandles")] | ||
| pub fn release_handles(keys: Vec<String>) { | ||
| HANDLE_TABLE.release_many(keys.into_iter().filter_map(|key| key.parse::<u64>().ok())); | ||
| } |
There was a problem hiding this comment.
Don’t silently ignore malformed keys in releaseHandles.
Current parsing drops invalid keys without signal, which can hide caller bugs and leak unreleased handles.
💡 Proposed hardening (non-breaking API)
#[wasm_bindgen(js_name = "releaseHandles")]
pub fn release_handles(keys: Vec<String>) {
- HANDLE_TABLE.release_many(keys.into_iter().filter_map(|key| key.parse::<u64>().ok()));
+ let mut parsed = Vec::with_capacity(keys.len());
+ for key in keys {
+ match key.parse::<u64>() {
+ Ok(v) => parsed.push(v),
+ Err(_) => log::warn!("Invalid handle key passed to releaseHandles: {key}"),
+ }
+ }
+ HANDLE_TABLE.release_many(parsed);
}
This PR makes opaque-handle cleanup predictable across runtimes (C/Python/WASM).
When we serialize an outbound value, we now also return
handles_created: the list of handle keys allocated during encoding. Hosts can reliably free any handles they don’t retain, preventing slow leaks in long-running processes.Cleanup is centralized with
HandleTable::release_manyplus shared binding helpers, replacing duplicated per-runtime logic and reducing lock contention and extra cloning.We also harden the FFI boundary by validating handle types returned by hosts (while staying forward-compatible with unknown enum values), and we avoid unexpectedly huge payloads by only inlining large media under a configurable size limit.
Migration:
BamlOutboundValuegainshandles_created, so consumers need to re-generate protobuf bindings to access the new field; host code should call the batch-release helper (client libs include helpers to compute/release unused handles).Summary by CodeRabbit
Release Notes