feat(middleware)!: make middleware async across primary bindings - #571
Merged
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThis PR changes middleware callbacks and lifecycle processing from synchronous to asynchronous across core APIs, Rust/Python/Node/FFI bindings, workers, native plugins, adaptive components, and PII redaction. It also adds queued FIFO event publication, async stream finalization, Native ABI v3 middleware support, and fail-open sanitizer handling. ChangesAsync middleware platform
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
This was referenced Jul 28, 2026
Member
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
2 tasks
There was a problem hiding this comment.
Actionable comments posted: 28
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
crates/plugin/tests/typed_callbacks.rs (1)
303-337: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAdd layout assertions for
NemoRelayNativeHostApiV3.The test now claims v3 coverage but only pins
NemoRelayNativeHostApiV1size/align/offsets. The v3 table is the newly frozen surface consumed by out-of-tree plugins; withoutsize_of/align_of/offset assertions (and a check thatoffset_of!(NemoRelayNativeHostApiV3, v1) == 0), a reordered or resized v3 field would silently break installed plugins that pass thestruct_sizegate inregister_async_middleware_raw.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/plugin/tests/typed_callbacks.rs` around lines 303 - 337, Add layout assertions for NemoRelayNativeHostApiV3 in native_abi_v3_struct_sizes_are_self_describing: verify its alignment and size, assert offset_of!(NemoRelayNativeHostApiV3, v1) == 0, and pin the complete field-offset array using the existing host_api_offsets-style helper or equivalent. Keep the existing V1, plugin, and stream assertions unchanged.crates/core/tests/unit/dynamic_worker_tests.rs (1)
1435-1484: 🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy liftFail-closed sanitizer fallback is still needed here
When the worker callback returns an error, these snapshot chains preserve the original payload (
event, tool args/result, and LLM request/response) instead of redacting or blocking emission. That leaks unsanitized observability data across mark, scope-start/end, tool request/response, and LLM request/response paths.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/core/tests/unit/dynamic_worker_tests.rs` around lines 1435 - 1484, Update the sanitizer snapshot-chain error handling exercised by this test so callback failures fail closed rather than returning the original payloads. Ensure event, tool request/response, and LLM request/response chains redact or block emission on errors, and revise these assertions to verify no unsanitized data is preserved.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/core/src/api/runtime/callbacks.rs`:
- Around line 30-37: Update EventSanitizeFn and event_sanitize_snapshot_chain to
share the event through Arc<Event> (or the minimal read-only context required)
rather than passing Event by value; adjust sanitizer implementations and call
sites so each sanitizer reuses the shared event without per-iteration deep
clones while preserving existing sanitization behavior.
In `@crates/core/src/api/runtime/state.rs`:
- Around line 995-1019: Async sanitizer failures currently restore unsanitized
values and continue, risking payload leaks. In
crates/core/src/api/runtime/state.rs:995-1019
(llm_sanitize_request_snapshot_chain), 1048-1072
(llm_sanitize_response_snapshot_chain), 684-703
(tool_sanitize_request_snapshot_chain), and 733-752
(tool_sanitize_response_snapshot_chain), choose and consistently implement an
explicit failure policy—preferably fail closed for sanitizers—rather than
restoring current; if retaining fail-open behavior, document it in the migration
guide and surface that policy consistently.
In `@crates/core/src/api/runtime/subscriber_dispatcher.rs`:
- Around line 220-222: Update the dispatcher loop’s DispatcherMessage::Barrier
handling to store barriers in a pending list instead of blocking on done.recv().
When processing Flush, wait for pending barriers to preserve ordering, using
recv_timeout and emitting a warning when a barrier times out or its producer is
dropped so the dispatcher continues. Ensure flush-from-finalizer paths cannot
deadlock the dispatcher thread.
In `@crates/core/src/api/scope.rs`:
- Around line 220-221: Update the documentation for push_scope in
crates/core/src/api/scope.rs:220-221 and event in
crates/core/src/api/scope.rs:351-352 to state that each event is queued for
dispatch with the sanitizer snapshot captured while the scope is active, rather
than promising subscriber observation before the function returns; no code
changes are required.
In `@crates/core/src/api/tool.rs`:
- Around line 284-303: Snapshot sanitizer failures currently drop events while
callers return success, orphaning spans. Apply one consistent fallback by
publishing the original unsanitized event when snapshot_event_sanitizers returns
None: update the tool START dispatch in crates/core/src/api/tool.rs lines
284-303, tool END dispatch in crates/core/src/api/tool.rs lines 466-489, and LLM
END dispatch in crates/core/src/api/llm.rs lines 876-878, preserving the
existing dispatch and completion behavior.
In `@crates/core/src/plugin/dynamic/native.rs`:
- Around line 1437-1464: Bound the await in the native callback flow around
NativeAsyncWait so a Pending callback cannot block indefinitely: wrap the
receiver wait with tokio::time::timeout, convert timeout expiry into
FlowError::Internal, and preserve the existing dropped-without-settling error
for receiver closure. Keep NativeAsyncWait ownership and drop-based cancellation
intact so expiry signals cancellation to native code.
- Around line 1594-1616: Update the NativeAsyncNextInner::LlmStream handling and
the analogous stream-interception path around the second occurrence so
downstream chunks are not buffered without bound before completing. Preserve
incremental delivery when the completion ABI supports it; otherwise enforce a
clear chunk-count or byte limit and return FlowError::Internal when exceeded,
documenting the enforced constraint.
In `@crates/core/src/plugin/dynamic/worker.rs`:
- Around line 1872-1887: Update the worker_callback_failed logging in the
invoke_async_with_timeout flow to log the surface name rather than the raw i32
surface value. Convert surface with RegistrationSurface::try_from(...).map(|s|
s.as_str_name()), matching the existing pattern in this file, while leaving the
callback behavior and returned result unchanged.
In `@crates/core/src/stream.rs`:
- Around line 344-355: Update the no-runtime branch in finish_with_status to use
the dedicated-thread finalization path already established by finish(), rather
than building a runtime and calling block_on inside poll_next. Preserve the
existing spawned-handle behavior when an ambient runtime is available, and
ensure finalizer middleware executes asynchronously without blocking the poller.
- Around line 437-451: Update the already-ended branch in the stream polling
logic to consume and return the stored terminal_result when finalization is
None, preserving the existing clean-end behavior when no terminal error is
stored. Ensure terminal errors assigned before finish_with_status are propagated
to the consumer instead of returning Poll::Ready(None).
In `@crates/core/tests/fixtures/native_plugin/src/lib.rs`:
- Around line 798-835: Update the async callback flow around the next.is_null()
|| completion.is_null() guard and subsequent reject branches to release the
non-null next handle before every failure return. Ensure the completion-null
case releases next, while retaining the existing successful-path release and
avoiding release when next itself is null.
In `@crates/ffi/src/api/mod.rs`:
- Around line 102-109: Document near the async registration API that synchronous
FFI middleware helpers must not be called from a Tokio runtime thread because
they return FlowError::Internal; direct embedded hosts to use the
completion-based async registration API instead. Update the surrounding API
documentation without changing block_on_sync_ffi behavior.
In `@crates/ffi/src/callable.rs`:
- Around line 706-712: Update wrap_llm_sanitize_request_fn and
wrap_llm_sanitize_response_fn to return the ffi_codec_identity error after
calling set_last_error, instead of returning Ok(None). In
crates/ffi/src/callable.rs lines 706-712 and 750-756, preserve error propagation
so sanitization failures cannot pass through as unchanged payloads; leave the
expect_err assertions in crates/ffi/tests/unit/callable_tests.rs lines 501-512
unchanged.
In `@crates/ffi/tests/integration/callable_extra_tests.rs`:
- Around line 7-18: Remove the duplicated resolve helper from the integration
test and move or reuse it from the shared FFI test-support module, consolidating
the current-thread Tokio blocking strategy. Update the affected tests to import
the shared resolve symbol, while preserving its existing generic future
behavior.
In `@crates/node/src/api/mod.rs`:
- Around line 3197-3208: Update the async wrapper flush_subscribers so it does
not call the blocking core_subscriber_api::flush_subscribers directly on a Tokio
worker; execute that wait via tokio::task::spawn_blocking, then propagate both
join and core errors through to_napi_err while preserving the existing Promise
behavior.
In `@crates/node/src/callable.rs`:
- Line 303: Update the promise-aware wrapper serialization paths around the
request handling at lines 303 and 379, plus the json construction near 406, to
propagate serialization errors instead of substituting Json::Null or panicking.
Match the explicit error mapping and recording behavior used by the rewritten
sync path around lines 747-753, ensuring the sanitizer receives the real
serialized request only after successful serialization.
- Around line 754-762: Replace the unnecessary payload clones in the callable
paths: pass the owned request directly to call_with_return_value, and move
fields.data and fields.metadata when constructing js_fields. Preserve the
existing callback behavior while consuming request and fields only after their
final use.
In `@crates/node/src/callback_factory.rs`:
- Around line 78-88: Update the rejection handler in the callback factory’s
promise chain to preserve non-string rejection reasons instead of leaving them
as “unknown error”. Safely stringify primitive and plain-object values, while
retaining the guarded error.message handling so throwing getters cannot escape;
align the fallback behavior with the sibling execution factory’s error
conversion.
In `@crates/node/src/promise_call.rs`:
- Around line 158-167: The reject callback in promise_call.rs expects a string,
but callback_factory.rs forwards the raw error on its early-return rejection
path. Update that wrapper branch to normalize the first error argument to a
string before calling reject, preserving the existing settle-path normalization
so both rejection routes deliver the actual error message.
In `@crates/node/tests/llm_tests.mjs`:
- Line 767: Strengthen the assertion following getLastCallbackError() so it
verifies the exact sanitizer error text, matching the response-side counterpart,
rather than accepting the broad “unknown error” or “callback” alternatives. Keep
the null-safe retrieval and ensure the test fails when the sanitizer exception
is not propagated correctly.
- Around line 60-69: Update waitForSubscriberCallbacks to call flushSubscribers
on each wait-loop iteration, before rechecking predicate, so callback batches
that become drainable later are processed without relying on native
self-draining; retain the existing timeout and polling behavior.
In `@crates/node/tests/tools_tests.mjs`:
- Around line 611-637: Update the TypeScript declarations in
registerToolConditionalExecutionGuardrail to accept string | null |
Promise<string | null>, and update registerToolRequestIntercept to accept Json |
Promise<Json>, matching the async callback behavior covered by the tests at
crates/node/tests/tools_tests.mjs lines 611-637 and 812-846; no direct test
changes are required.
In `@crates/pii-redaction/src/builtin.rs`:
- Around line 464-472: Wrap the compiled backend in an Arc once when
constructing the callbacks, then clone the Arc inside each invocation instead of
cloning CompiledBuiltinBackend. Apply this consistently to the callbacks around
the shown closure and to event_sanitize_callback_with_scope_categories,
llm_sanitize_request_callback, and llm_sanitize_response_callback, preserving
their existing async sanitization behavior.
In `@crates/python/src/py_api/mod.rs`:
- Around line 1324-1342: Update the synchronous branch around
`tool_request_intercepts` to reuse `pyo3_async_runtimes::tokio::get_runtime()`
instead of constructing a new `tokio::runtime::Builder` per invocation. Apply
the same change across all four helpers, preserving the existing
`TASK_SCOPE_STACK` scoping, error mapping, and result conversion; optionally
centralize the repeated sync-or-future branching in a `run_sync_or_future`
helper.
- Around line 1320-1340: Update the synchronous branch using runtime.block_on
and the awaitable handling in split_py_object_or_future/split_json_or_future so
sync callers cannot attempt to execute coroutine returns through
pyo3_async_runtimes::tokio::into_future. Detect awaitable middleware in this
path and return a clear Python-facing error, while preserving awaitable support
for async callers.
In `@crates/python/src/py_callable.rs`:
- Around line 1119-1145: Remove the eprintln! calls from the three
conversion-error branches handling py_event, fields_json, and py_fields.
Preserve each branch’s existing FlowError::Internal(error.to_string()) return so
failures are propagated without duplicate stderr reporting.
In `@crates/python/tests/coverage/py_api_coverage_tests.rs`:
- Around line 475-499: Add async coverage in
crates/python/tests/coverage/py_api_coverage_tests.rs:475-499 by using
with_event_loop and run_until_complete to await results from
tool_request_intercepts, tool_conditional_execution, llm_request_intercepts, and
llm_conditional_execution, including success and error behavior. In
crates/python/tests/coverage/coverage_tests.rs:664-681, define module functions
with async def so the wrappers exercise split_*_or_future, and assert both
successful results and awaited failures.
In `@crates/python/tests/coverage/py_callable_coverage_tests.rs`:
- Around line 703-728: Rename the test function
event_sanitize_wrapper_covers_conversion_success_and_fail_closed_paths to
describe conversion success and error propagation, without changing its
assertions or behavior. Keep the test focused on the wrapper returning errors
for raised and invalid sanitizer results.
---
Outside diff comments:
In `@crates/core/tests/unit/dynamic_worker_tests.rs`:
- Around line 1435-1484: Update the sanitizer snapshot-chain error handling
exercised by this test so callback failures fail closed rather than returning
the original payloads. Ensure event, tool request/response, and LLM
request/response chains redact or block emission on errors, and revise these
assertions to verify no unsanitized data is preserved.
In `@crates/plugin/tests/typed_callbacks.rs`:
- Around line 303-337: Add layout assertions for NemoRelayNativeHostApiV3 in
native_abi_v3_struct_sizes_are_self_describing: verify its alignment and size,
assert offset_of!(NemoRelayNativeHostApiV3, v1) == 0, and pin the complete
field-offset array using the existing host_api_offsets-style helper or
equivalent. Keep the existing V1, plugin, and stream assertions unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Enterprise
Run ID: c5ddd3c2-e095-4e58-b84f-292045630fe3
📒 Files selected for processing (77)
crates/adaptive/src/acg_component.rscrates/adaptive/src/adaptive_hints_intercept.rscrates/adaptive/src/lib.rscrates/adaptive/tests/integration/runtime_integration_tests.rscrates/adaptive/tests/support/mod.rscrates/adaptive/tests/unit/acg_component_tests.rscrates/adaptive/tests/unit/adaptive_hints_intercept_tests.rscrates/adaptive/tests/unit/plugin_component_tests.rscrates/adaptive/tests/unit/runtime_features_tests.rscrates/adaptive/tests/unit/runtime_tests.rscrates/cli/src/sessions/mod.rscrates/cli/tests/coverage/shared/server_tests.rscrates/core/src/api/llm.rscrates/core/src/api/runtime/callbacks.rscrates/core/src/api/runtime/state.rscrates/core/src/api/runtime/subscriber_dispatcher.rscrates/core/src/api/scope.rscrates/core/src/api/shared.rscrates/core/src/api/tool.rscrates/core/src/logging/rotation.rscrates/core/src/plugin/dynamic/native.rscrates/core/src/plugin/dynamic/worker.rscrates/core/src/stream.rscrates/core/tests/coverage/logging_rotation_tests.rscrates/core/tests/coverage/logging_sink_tests.rscrates/core/tests/fixtures/native_plugin/src/lib.rscrates/core/tests/integration/api_surface_tests.rscrates/core/tests/integration/middleware_tests.rscrates/core/tests/integration/native_plugin_tests.rscrates/core/tests/integration/pipeline_tests.rscrates/core/tests/integration/scope_local_tests.rscrates/core/tests/integration/subscriber_dispatcher_tests.rscrates/core/tests/integration/test_support.rscrates/core/tests/integration/worker_plugin_tests.rscrates/core/tests/unit/context_tests.rscrates/core/tests/unit/dynamic_worker_tests.rscrates/core/tests/unit/llm_api_tests.rscrates/core/tests/unit/native_plugin_tests.rscrates/core/tests/unit/plugin_tests.rscrates/core/tests/unit/shared_tests.rscrates/ffi/src/api/mod.rscrates/ffi/src/callable.rscrates/ffi/tests/integration/callable_extra_tests.rscrates/ffi/tests/unit/callable_tests.rscrates/node/src/api/mod.rscrates/node/src/callable.rscrates/node/src/callback_factory.rscrates/node/src/promise_call.rscrates/node/tests/callback_error_tests.mjscrates/node/tests/event_sanitizers_tests.mjscrates/node/tests/llm_tests.mjscrates/node/tests/scope_tests.mjscrates/node/tests/tools_tests.mjscrates/pii-redaction/src/builtin.rscrates/pii-redaction/tests/unit/component_tests.rscrates/plugin/README.mdcrates/plugin/src/lib.rscrates/plugin/tests/typed_callbacks.rscrates/python/src/py_api/mod.rscrates/python/src/py_callable.rscrates/python/tests/coverage/coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_callable_coverage_tests.rsdocs/about-nemo-relay/concepts/middleware.mdxdocs/build-plugins/dynamic-plugins/native-dynamic/about.mdxdocs/reference/event-sanitizers.mdxdocs/reference/migration-guides.mdxintegrations/openclaw/test/live-smoke.test.tspython/nemo_relay/__init__.pypython/nemo_relay/__init__.pyipython/nemo_relay/_native.pyipython/tests/test_adaptive.pypython/tests/test_builtin_codecs.pypython/tests/test_context_isolation.pypython/tests/test_event_sanitizers.pypython/tests/test_llm.pypython/tests/test_tools.py
willkill07
force-pushed
the
wkk_relay-509-async-primary-middleware
branch
from
July 28, 2026 18:41
7934489 to
9be032b
Compare
willkill07
force-pushed
the
wkk_relay-509-event-publication
branch
from
July 28, 2026 18:44
2ba7152 to
d5a95dc
Compare
willkill07
force-pushed
the
wkk_relay-509-async-primary-middleware
branch
from
July 28, 2026 18:44
9be032b to
054ce5e
Compare
willkill07
force-pushed
the
wkk_relay-509-async-primary-middleware
branch
from
July 28, 2026 19:34
7e07801 to
7a5a0ef
Compare
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
bbednarski9
reviewed
Jul 29, 2026
bbednarski9
left a comment
Contributor
There was a problem hiding this comment.
I reproduced four async-callback correctness failures against this exact PR head with focused smoke tests. The corresponding awaited/non-spawned control cases pass, so these are specific lifetime, context, concurrency, and publication-ordering regressions rather than harness failures.
Salonijain27
approved these changes
Jul 30, 2026
Salonijain27
left a comment
Collaborator
There was a problem hiding this comment.
Approved from a dependency point of view
2 tasks
2 tasks
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
bbednarski9
reviewed
Jul 30, 2026
bbednarski9
left a comment
Contributor
There was a problem hiding this comment.
One new async-continuation finding, reproduced with focused Node tool and unary-LLM smoke tests and isolated against passing Rust/Python controls.
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <wkillian@nvidia.com>
bbednarski9
reviewed
Jul 30, 2026
bbednarski9
left a comment
Contributor
There was a problem hiding this comment.
I re-ran focused smoke tests and controls against the current 698a5f6 head. These three defects remain reproducible across the affected Rust, Python, and Node paths; the paired normal-close, explicit-scope, and intercepted-stream controls pass.
Signed-off-by: Will Killian <wkillian@nvidia.com>
Signed-off-by: Will Killian <2007799+willkill07@users.noreply.github.com>
bbednarski9
reviewed
Jul 30, 2026
bbednarski9
left a comment
Contributor
There was a problem hiding this comment.
Re-reviewed current head 9bdb5c9 with a rebuilt Node addon. The previously reported terminal cleanup, default stream context, sibling branch isolation, and typed cancellation regressions now pass. These three additional Node context/lifetime defects remain directly reproducible; a fourth follow-up is attached to the existing scope-stack thread.
bbednarski9
approved these changes
Jul 30, 2026
bbednarski9
left a comment
Contributor
There was a problem hiding this comment.
approving w/ understanding we will disclaim known Node failures and patch those later.
Member
Author
|
/merge |
This was referenced Jul 31, 2026
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
4 participants
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.
Overview
Convert middleware callbacks to an async Rust contract and add awaitable/Promise support across the primary Python and Node.js bindings.
Warning
BREAKING CHANGE: Rust middleware callbacks now return futures. Update guardrails, request/execution intercepts, tool/LLM sanitizers, and event sanitizers to return
Box::pin(async move { ... })(or an equivalent future). Native plugins must be rebuilt against context ABI v3. Raw C FFI and Go middleware callbacks remain synchronous and occupy a native thread while running. Python and Node.js support awaitables/Promises additively; scope, mark, and manual lifecycle event-emission APIs remain synchronous.Details
This is PR 2 of 2 in the RELAY-509 GitHub stack. The stack is rooted on upstream
main, and this PR is based on #570:Where should the reviewer start?
Start with
crates/core/src/api/runtime/callbacks.rs, then review the Python adapter incrates/python/src/py_callable.rsand Node bridge incrates/node/src/promise_call.rs.Validation:
just test-rust— passedjust test-python— 595 passedjust test-node— 325 passedjust test-go— passed; raw C FFI middleware remains synchronousjust docsanduv run pre-commit run --all-files— passedRelated Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
New Features
Bug Fixes
Documentation