bex_engine: add CancellationToken support across all layers#3136
Merged
antoniosarosi merged 10 commits intocanaryfrom Feb 19, 2026
Merged
bex_engine: add CancellationToken support across all layers#3136antoniosarosi merged 10 commits intocanaryfrom
antoniosarosi merged 10 commits intocanaryfrom
Conversation
Implement cooperative cancellation for BAML function execution: - Phase 1: Core CancellationToken type, biased select in VM event loop, AbortHandle for killing spawned tasks, EngineError::Cancelled variant - Phase 2: Wire cancellation through SysOpContext so sys_ops (sleep, HTTP) respect the token; add is_cancelled() builtin - Phase 3: Cooperative cancellation checks in llm.baml orchestrator (retry loop, primitive/fallback branches, top-level call) - Phase 4: Language binding integration - Python AbortController, WASM callId-based cancellation, C FFI pass-through - Tests: 8 integration tests covering immediate cancel, sleep/HTTP interruption, selective cancellation, idempotency, and normal completion
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
|
Important Review skippedReview was skipped as selected files did not have any reviewable changes. 💤 Files selected but had no reviewable changes (4)
⛔ Files ignored due to path filters (32)
You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds workspace dependency Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Orchestrator as LLM_Orchestrator
participant Executor as execute_client
participant SysOp as cancellation_requested
Client->>Orchestrator: call_llm_function(...)
Orchestrator->>Executor: execute_client(...)
Executor->>SysOp: cancellation_requested()
SysOp-->>Executor: true/false
alt cancellation == true
Executor-->>Orchestrator: short-circuit / cancel
Orchestrator-->>Client: propagate cancellation/error
else cancellation == false
Executor-->>Orchestrator: continue (attempt / primitive / fallback)
Orchestrator-->>Client: result or final failure
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Clearer name — makes it obvious this checks whether cancellation has been requested for the current call, not some ambient state.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
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/baml_builtins/baml/llm.baml (1)
263-270:⚠️ Potential issue | 🟡 Minor
RoundRobinbranch advances the atomic counter before checking cancellation.
Fallbackcheckscancellation_requested()before each sub-client iteration;RoundRobincallsround_robin_next(advancing the counter atomically) before any cancellation check. If the token is already set, the counter slot is wasted and the client selection is skewed. The actual call is still caught by the checks insideexecute_client/execute_client_once, but the counter side-effect has already occurred.🛡️ Proposed fix
baml.llm.ClientType.RoundRobin => { + if (baml.sys.cancellation_requested()) { + return baml.llm.ExecutionResult { ok: false, value: null }; + } let idx = baml.llm.round_robin_next(llm_client.name) % llm_client.sub_clients.length(); baml.llm.execute_client( llm_client.sub_clients.at(idx),
- Add early cancellation check at top of call_function so pre-cancelled tokens always produce Err(Cancelled) regardless of function contents - Upgrade any error to Cancelled when the token is cancelled, so cooperative BAML-level checks (via baml.sys.panic) are reported as Cancelled to callers - Fix poisoned mutex handling in C FFI (unwrap_or_else instead of silent skip) - Error on WASM call_id collision instead of silently replacing - Document WASM orphan future limitation and cooperative check semantics - Split cancellation_requested test into false/true cases, remove stale comment
true is only observable if all preceding sys_ops were sync; any async op would trigger VM-level cancellation at its Await first.
Replace platform-specific abort logic with futures::future::AbortHandle which works on both native (tokio::spawn) and WASM (spawn_local). WASM orphan futures are now properly aborted on cancellation instead of running to completion.
Contributor
There was a problem hiding this comment.
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/baml_builtins/baml/llm.baml (1)
268-276: 🧹 Nitpick | 🔵 TrivialNo cancellation check in the RoundRobin branch of
execute_client_once.The Primitive (Line 222) and Fallback (Line 253) branches both have early
cancellation_requested()guards, but the RoundRobin branch at Line 268 does not. This is likely acceptable because:
round_robin_nextis an async sys_op, so the VM's biasedselect!would catch cancellation at that await point.- The recursive
execute_clientcall will hit the cancellation check in the resolved child's branch.However, the
round_robin_nextcall still executes and mutates the atomic counter even when cancellation is already requested. If counter accuracy across cancelled calls matters, consider adding a guard here too.Optional: add cancellation guard before round-robin counter mutation
baml.llm.ClientType.RoundRobin => { + if (baml.sys.cancellation_requested()) { + return baml.llm.ExecutionResult { ok: false, value: null }; + } let idx = baml.llm.round_robin_next(llm_client.name) % llm_client.sub_clients.length();
Prevents round_robin_next from mutating the atomic counter when cancellation is already requested, matching the Primitive and Fallback branches.
C FFI: reject duplicate call_ids with DuplicateCallId error instead of silently overwriting the previous entry (matching WASM behavior). WASM: surface EngineError::Cancelled as a JS Error with name="BamlCancelledError" so callers can programmatically distinguish cancellation from other failures (matching Python's BamlCancelledError).
# Conflicts: # baml_language/crates/baml_tests/snapshots/attribute_validation/baml_tests__attribute_validation__06_codegen.snap # baml_language/crates/baml_tests/snapshots/basic_types/baml_tests__basic_types__06_codegen.snap # baml_language/crates/baml_tests/snapshots/builtin_io/baml_tests__builtin_io__06_codegen.snap # baml_language/crates/baml_tests/snapshots/comment_after_string_in_config/baml_tests__comment_after_string_in_config__06_codegen.snap # baml_language/crates/baml_tests/snapshots/comment_in_type/baml_tests__comment_in_type__06_codegen.snap # baml_language/crates/baml_tests/snapshots/config_dictionary/baml_tests__config_dictionary__06_codegen.snap # baml_language/crates/baml_tests/snapshots/duplicate_class_span/baml_tests__duplicate_class_span__06_codegen.snap # baml_language/crates/baml_tests/snapshots/function_call/baml_tests__function_call__06_codegen.snap # baml_language/crates/baml_tests/snapshots/function_types/baml_tests__function_types__06_codegen.snap # baml_language/crates/baml_tests/snapshots/generator/baml_tests__generator__06_codegen.snap # baml_language/crates/baml_tests/snapshots/header_in_llm_function/baml_tests__header_in_llm_function__06_codegen.snap # baml_language/crates/baml_tests/snapshots/headers_edge_cases/baml_tests__headers_edge_cases__06_codegen.snap # baml_language/crates/baml_tests/snapshots/match_exhaustiveness/baml_tests__match_exhaustiveness__06_codegen.snap # baml_language/crates/baml_tests/snapshots/paren_union_test/baml_tests__paren_union_test__06_codegen.snap # baml_language/crates/baml_tests/snapshots/parser_constructors/baml_tests__parser_constructors__06_codegen.snap # baml_language/crates/baml_tests/snapshots/parser_error_recovery/baml_tests__parser_error_recovery__06_codegen.snap # baml_language/crates/baml_tests/snapshots/parser_expressions/baml_tests__parser_expressions__06_codegen.snap # baml_language/crates/baml_tests/snapshots/parser_speculative/baml_tests__parser_speculative__06_codegen.snap # baml_language/crates/baml_tests/snapshots/parser_statements/baml_tests__parser_statements__06_codegen.snap # baml_language/crates/baml_tests/snapshots/parser_stress/baml_tests__parser_stress__06_codegen.snap # baml_language/crates/baml_tests/snapshots/parser_strings/baml_tests__parser_strings__06_codegen.snap # baml_language/crates/baml_tests/snapshots/pending_greaters_fix/baml_tests__pending_greaters_fix__06_codegen.snap # baml_language/crates/baml_tests/snapshots/retry_policy/baml_tests__retry_policy__06_codegen.snap # baml_language/crates/baml_tests/snapshots/simple_function/baml_tests__simple_function__06_codegen.snap # baml_language/crates/baml_tests/snapshots/simple_type_error/baml_tests__simple_type_error__06_codegen.snap # baml_language/crates/baml_tests/snapshots/top_level_header_comment/baml_tests__top_level_header_comment__06_codegen.snap # baml_language/crates/baml_tests/snapshots/top_level_let/baml_tests__top_level_let__06_codegen.snap # baml_language/crates/baml_tests/snapshots/type_aliases/baml_tests__type_aliases__06_codegen.snap # baml_language/crates/baml_tests/snapshots/type_builder_errors/baml_tests__type_builder_errors__06_codegen.snap # baml_language/crates/baml_tests/snapshots/type_builder_test/baml_tests__type_builder_test__06_codegen.snap # baml_language/crates/baml_tests/snapshots/unknown_type_error/baml_tests__unknown_type_error__06_codegen.snap
Merging this PR will degrade performance by 11.63%
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ❌ | WallTime | bench_single_simple_file |
926.7 µs | 1,046.2 µs | -11.43% |
| ❌ | WallTime | bench_incremental_add_string_char |
918.2 µs | 1,039.1 µs | -11.63% |
| ❌ | WallTime | bench_incremental_close_string |
921.2 µs | 1,038.6 µs | -11.3% |
| ❌ | WallTime | bench_incremental_add_attribute |
924.1 µs | 1,037.1 µs | -10.9% |
| ❌ | WallTime | bench_empty_project |
855.5 µs | 963.3 µs | -11.19% |
Comparing antonio/cancellation (64208f9) with canary (8838648)
Footnotes
-
84 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩
Binary size checks passed✅ 7 passed
Generated by |
tokio::select! requires the macros feature, which was only enabled for native targets. The select! macro is a pure proc macro with no runtime dependency, so it works on WASM.
This was referenced Feb 20, 2026
Merged
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
CancellationToken(wrappingtokio_util::sync::CancellationToken), addEngineError::Cancelledvariant, integrate biasedtokio::select!in the VM event loop so cancellation is checked before every await, and useAbortHandleto kill spawned tasks on cancel.CancellationTokenthroughSysOpContextso all sys_ops (sleep, HTTP fetch, etc.) respect cancellation. Addbaml.sys.is_cancelled()builtin for cooperative polling.is_cancelled()checks inllm.bamlat retry loops, primitive/fallback branches, and the top-levelcall_llm_functionentry point.AbortControllerclass withabort()/abortedAPI, WASMcallFunction(name, args, callId?)+cancelCall(callId)pattern, C FFI pass-through (already wired in Phase 1).bex_engine/tests/cancellation.rscovering immediate cancel, sleep/HTTP interruption, selective cancellation, sequential sleeps, idempotent cancel, cooperativeis_cancelled(), and normal completion.Test plan
cargo test -p bex_engine --test cancellation— all 8 tests passcargo clippy --workspace— cleanis_cancelledglobal index shift)Summary by CodeRabbit
New Features
Breaking Changes
Chores