Skip to content

bridge(feat): track created handles and enable batch cleanup#3195

Open
miguelcsx wants to merge 1 commit intoBoundaryML:canaryfrom
miguelcsx:feat/handle-type
Open

bridge(feat): track created handles and enable batch cleanup#3195
miguelcsx wants to merge 1 commit intoBoundaryML:canaryfrom
miguelcsx:feat/handle-type

Conversation

@miguelcsx
Copy link
Copy Markdown
Contributor

@miguelcsx miguelcsx commented Mar 1, 2026

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_many plus 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: BamlOutboundValue gains handles_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

  • New Features
    • Added bulk handle release capability across all language bindings.
    • Implemented automatic cleanup of unused handles after operations.
    • Added handle type validation to ensure type correctness.
    • Introduced configurable size limits for inlined media content.

- 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
@vercel
Copy link
Copy Markdown

vercel bot commented Mar 1, 2026

@miguelcsx is attempting to deploy a commit to the Boundary Team on Vercel.

A member of the Team first needs to authorize it.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Mar 1, 2026

📝 Walkthrough

Walkthrough

This 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 release_many method to release multiple handles in a single call. Protobuf definitions and generated files are updated to support handle tracking.

Changes

Cohort / File(s) Summary
Bulk Handle Release API
baml_language/crates/bridge_cffi/src/ffi/handle.rs, baml_language/crates/bridge_cffi/src/lib.rs, baml_language/crates/bridge_python/src/handle.rs, baml_language/crates/bridge_wasm/src/handle.rs, baml_language/crates/bridge_wasm/src/lib.rs
Adds release_handles/release_many FFI entry points across CFFI, Python, and WASM bridges to release multiple handles in a single call.
HandleTable Core Implementation
baml_language/crates/bridge_ctypes/src/handle_table.rs
Introduces release_many() method for bulk handle release, adds max_inline_media_bytes configuration option to HandleTableOptions, and implements TryFrom<&BexExternalValue> for HandleTableValue.
Handle Type Validation
baml_language/crates/bridge_ctypes/src/error.rs, baml_language/crates/bridge_ctypes/src/value_decode.rs
Adds InvalidHandleType error variant and introduces validate_handle_type() to enforce type consistency during inbound handle decoding.
Outbound Value Handle Tracking
baml_language/crates/bridge_ctypes/src/value_encode.rs
Refactors encoding pipeline to collect and track handles created during external-to-BAML conversion, implements media inlining size limits via should_inline_media() helper, and updates return flow to include handles in BamlOutboundValue.
Proto Definitions & Generated Code
baml_language/crates/bridge_ctypes/types/baml/cffi/v1/baml_outbound.proto, baml_language/crates/bridge_python/python_src/baml/cffi/v1/baml_outbound_pb2.py, baml_language/crates/bridge_python/python_src/baml/cffi/v1/baml_inbound_pb2.py
Adds handles_created repeated field to BamlOutboundValue proto; regenerates Python protobuf descriptors with updated serialization offsets.
Python Handle Cleanup
baml_language/crates/bridge_python/python_src/baml_py/proto.py
Adds _collect_handles() helper to traverse and collect used handle keys, extends decode_call_result() to release unused handles via bulk release API.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Possibly related PRs

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'bridge(feat): track created handles and enable batch cleanup' accurately and concisely captures the main objective of the PR: implementing handle tracking and batch cleanup across runtimes.
Docstring Coverage ✅ Passed Docstring coverage is 82.05% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 | 🟠 Major

PromptAst media bypasses max_inline_media_bytes and can still inline oversized payloads.

When serialize_prompt_ast is enabled, media inside PromptAstSimple::Media is always inlined via bex_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

📥 Commits

Reviewing files that changed from the base of the PR and between ac50196 and b2f890c.

📒 Files selected for processing (13)
  • baml_language/crates/bridge_cffi/src/ffi/handle.rs
  • baml_language/crates/bridge_cffi/src/lib.rs
  • baml_language/crates/bridge_ctypes/src/error.rs
  • baml_language/crates/bridge_ctypes/src/handle_table.rs
  • baml_language/crates/bridge_ctypes/src/value_decode.rs
  • baml_language/crates/bridge_ctypes/src/value_encode.rs
  • baml_language/crates/bridge_ctypes/types/baml/cffi/v1/baml_outbound.proto
  • baml_language/crates/bridge_python/python_src/baml/cffi/v1/baml_inbound_pb2.py
  • baml_language/crates/bridge_python/python_src/baml/cffi/v1/baml_outbound_pb2.py
  • baml_language/crates/bridge_python/python_src/baml_py/proto.py
  • baml_language/crates/bridge_python/src/handle.rs
  • baml_language/crates/bridge_wasm/src/handle.rs
  • baml_language/crates/bridge_wasm/src/lib.rs

Comment on lines +209 to +221
/// 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);
}
}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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."

Comment on lines +196 to +208
#[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"),
}
}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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."

Comment on lines +26 to +29
#[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()));
}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

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);
 }

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant