feat: add in-process Rampart PII redaction plugin - #558
Conversation
WalkthroughAdds an in-process Rampart PII plugin (detection, redaction, host registration, CLI editing, Node/Go/Python APIs) and converts NeMo Relay's middleware pipeline—guardrails, sanitizers, intercepts, and streaming execution—from synchronous to asynchronous callbacks across core, native ABI (v2→v3), FFI, Node, Python, worker, and adaptive crates, with matching tests and documentation. ChangesRampart PII plugin
Async middleware refactor
Estimated code review effort: 5 (Critical) | ~180 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
size:XXL
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
b73ac66 to
385c241
Compare
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
afourniernv
changed the title
afourniernv
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/worker/src/lib.rs (1)
676-705: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMerge
push_registrationintopush_contract_registrationto remove duplication.Both functions build an identical
Registrationexcept for thecontractfield. Python's SDK already unifies this via a single_push_registration(..., *, contract: str = "")(seepython/plugin/src/nemo_relay_plugin/_api.pylines 1288-1310); the Rust SDK diverging into two near-duplicate helpers risks future drift (e.g., a new field added to one but not the other).♻️ Proposed refactor
- fn push_registration( - &mut self, - name: &str, - surface: RegistrationSurface, - priority: i32, - break_chain: bool, - ) { - self.handlers.registrations.push(Registration { - local_name: name.into(), - surface: surface as i32, - priority, - break_chain, - contract: String::new(), - }); - } - - fn push_contract_registration( - &mut self, - name: &str, - surface: RegistrationSurface, - contract: &str, - ) { - self.handlers.registrations.push(Registration { - local_name: name.into(), - surface: surface as i32, - priority: 0, - break_chain: false, - contract: contract.into(), - }); - } + fn push_registration( + &mut self, + name: &str, + surface: RegistrationSurface, + priority: i32, + break_chain: bool, + ) { + self.push_contract_registration(name, surface, priority, break_chain, ""); + } + + fn push_contract_registration( + &mut self, + name: &str, + surface: RegistrationSurface, + priority: i32, + break_chain: bool, + contract: &str, + ) { + self.handlers.registrations.push(Registration { + local_name: name.into(), + surface: surface as i32, + priority, + break_chain, + contract: contract.into(), + }); + }Then update the call in
register_worker_inferencetoself.push_contract_registration(name, RegistrationSurface::WorkerInference, 0, false, contract);.🤖 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/worker/src/lib.rs` around lines 676 - 705, Merge push_registration into push_contract_registration by giving push_contract_registration priority and break_chain parameters, with contract supplied as the final argument and defaulting to an empty string where appropriate. Remove the duplicate helper, update all callers including register_worker_inference to pass the unified arguments, and preserve existing Registration field values.
🤖 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/pii-redaction/src/component.rs`:
- Around line 960-962: Update is_valid_json_pointer_pattern to validate wildcard
segments according to JsonPointerPattern::matches: reject any path segment that
contains '*' unless the entire segment is exactly "*". Preserve existing JSON
Pointer validation for all other segments and continue accepting standalone
wildcard segments.
In `@crates/pii-redaction/src/local.rs`:
- Around line 732-740: The empty-paths branch in llm_sanitize_request_callback
must sanitize request headers as well as request.content, preserving the
intended pointer prefixes for the headers and content roots. Reuse the existing
request/header sanitization behavior used by sanitize_raw_request or the builtin
flow, and ensure the broad-coverage path returns a request with sanitized header
values.
In `@crates/pii-redaction/tests/unit/component_tests.rs`:
- Around line 1916-1921: Update the table-driven assertion in the
validate_plugin_config test loop to include failure context identifying the
current config/field/message case and the produced diagnostics. Preserve the
existing matching condition while supplying a descriptive assertion message so
failures reveal which case failed and the actual report contents.
In `@crates/pii-redaction/tests/worker_detection_tests.rs`:
- Around line 54-58: Remove the duplicate “/message” selector from either
target_paths or target_path_patterns in the test configuration, keeping it in
only one collection so the test expresses a single intent.
- Around line 268-285: Update the fail-closed assertion in the worker exit test
around the event emitted by “worker-pii-exit” to use a message value that the
healthy fixture worker does not redact, while retaining the expected redaction
for “unselected” if applicable. Ensure the assertion can only pass when the
crashed batch is handled fail-closed, rather than matching normal “PRIVATE”
detection behavior.
In `@crates/pii-redaction/workers/rampart/README.md`:
- Around line 134-136: Update the Runtime Bounds section in the README by adding
a complete introductory sentence before the existing bullet list; leave the
documented limits unchanged and ensure the lead-in grammatically introduces the
list.
In `@docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx`:
- Line 168: Rename the “Register worker inference” heading to “Register Worker
Inference” in docs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdx at
lines 168-168 and docs/build-plugins/dynamic-plugins/grpc-worker/rust/about.mdx
at lines 71-71, preserving the existing heading structure.
In `@docs/configure-plugins/pii-redaction/configuration.mdx`:
- Around line 424-431: The omission guidance following the “Path Semantics”
section is outdated. Update the paragraph describing manual LLM calls with
normalized target_paths and no active or fallback codec to state that payloads
are sanitized using the configured raw paths and emitted, matching the
early-return behavior in builtin and local redaction flows; preserve the
documented fail-closed contract.
- Around line 329-331: Add a complete introductory sentence immediately before
the TOML code block following the sanitizer registration-rejection paragraph,
clearly describing what the configuration example demonstrates. Keep the
existing TOML content unchanged.
In `@go/nemo_relay/pii_redaction/pii_redaction_test.go`:
- Around line 47-62: Extend the validation condition in the NewComponentSpec
test to assert that spec.Config.Local.Backend matches the configured backend
value from the test setup. Keep the existing configuration assertions unchanged
and include the backend check alongside the other Local fields.
In `@python/plugin/README.md`:
- Around line 108-125: Update the Worker Inference example to establish that ctx
is a PluginContext available inside WorkerPlugin.register, either by showing the
enclosing register method or explicitly stating that scope. Keep the
register_worker_inference usage and handler behavior unchanged.
---
Outside diff comments:
In `@crates/worker/src/lib.rs`:
- Around line 676-705: Merge push_registration into push_contract_registration
by giving push_contract_registration priority and break_chain parameters, with
contract supplied as the final argument and defaulting to an empty string where
appropriate. Remove the duplicate helper, update all callers including
register_worker_inference to pass the unified arguments, and preserve existing
Registration field values.
🪄 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: b6f7fc1b-5518-49b8-af16-dd1759b6fa64
📒 Files selected for processing (60)
crates/cli/src/server/mod.rscrates/core/src/lib.rscrates/core/src/plugin.rscrates/core/src/plugin/dynamic/host.rscrates/core/src/plugin/dynamic/worker.rscrates/core/src/plugin/worker_inference.rscrates/core/tests/fixtures/worker_plugin/src/main.rscrates/core/tests/integration/worker_plugin_tests.rscrates/core/tests/unit/dynamic_worker_tests.rscrates/core/tests/unit/plugin_tests.rscrates/core/tests/unit/worker_inference_tests.rscrates/node/pii_redaction.d.tscrates/node/pii_redaction.jscrates/node/tests/pii_redaction_tests.mjscrates/pii-redaction/Cargo.tomlcrates/pii-redaction/README.mdcrates/pii-redaction/src/builtin.rscrates/pii-redaction/src/component.rscrates/pii-redaction/src/local.rscrates/pii-redaction/tests/unit/component_tests.rscrates/pii-redaction/tests/unit/local_tests.rscrates/pii-redaction/tests/worker_detection_tests.rscrates/pii-redaction/workers/rampart/MANIFEST.incrates/pii-redaction/workers/rampart/README.mdcrates/pii-redaction/workers/rampart/THIRD_PARTY_NOTICES.mdcrates/pii-redaction/workers/rampart/config.schema.jsoncrates/pii-redaction/workers/rampart/nemo_relay_pii_rampart/__init__.pycrates/pii-redaction/workers/rampart/nemo_relay_pii_rampart/detector.pycrates/pii-redaction/workers/rampart/nemo_relay_pii_rampart/prefetch.pycrates/pii-redaction/workers/rampart/nemo_relay_pii_rampart/py.typedcrates/pii-redaction/workers/rampart/nemo_relay_pii_rampart/worker.pycrates/pii-redaction/workers/rampart/pyproject.tomlcrates/pii-redaction/workers/rampart/relay-plugin.tomlcrates/pii-redaction/workers/rampart/tests/test_detector.pycrates/pii-redaction/workers/rampart/tests/test_worker.pycrates/worker-proto/README.mdcrates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker.protocrates/worker-proto/tests/proto_tests.rscrates/worker/README.mdcrates/worker/src/lib.rscrates/worker/tests/worker_sdk_tests.rsdocs/about-nemo-relay/release-notes/index.mdxdocs/build-plugins/dynamic-plugins/grpc-worker/grpc-worker-protocol.mdxdocs/build-plugins/dynamic-plugins/grpc-worker/python/about.mdxdocs/build-plugins/dynamic-plugins/grpc-worker/rust/about.mdxdocs/configure-plugins/pii-redaction/about.mdxdocs/configure-plugins/pii-redaction/configuration.mdxgo/nemo_relay/pii_redaction.gogo/nemo_relay/pii_redaction/pii_redaction.gogo/nemo_relay/pii_redaction/pii_redaction_test.gogo/nemo_relay/pii_redaction_test.gojustfilepython/nemo_relay/pii_redaction.pypython/nemo_relay/pii_redaction.pyipython/plugin/README.mdpython/plugin/src/nemo_relay_plugin/__init__.pypython/plugin/src/nemo_relay_plugin/_api.pypython/tests/plugin/test_public_api_docstrings.pypython/tests/plugin/test_worker_sdk.pypython/tests/test_pii_redaction_plugin.py
Signed-off-by: Alex Fournier <afournier@nvidia.com>
There was a problem hiding this comment.
You could refine local_model to resolve an integration-provided callback. The PII middleware would execute that callback with the selected text and detector settings, then receive the detected spans, labels, and confidence scores. The integration could implement the callback using any model, runtime, or transport it chooses, while the PII component continues to own field selection, detection validation, policy, and redaction.
github-actions
Bot
added
lang:go
There was a problem hiding this comment.
Actionable comments posted: 64
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
crates/core/src/api/runtime/scope_stack.rs (1)
446-458: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winNew public API can panic on a poisoned scope-stack lock.
task_scope_top()usesexpect("scope stack lock poisoned"), socapture_propagation_context*panics instead of returning theResultit already declares. This PR converts the same poison case toFlowError::Internalincrates/core/src/api/scope.rs(scope_stack_lock_error); reading the top handle here withoutexpectkeeps the boundary-capture API consistent and non-panicking.🤖 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/src/api/runtime/scope_stack.rs` around lines 446 - 458, Update capture_propagation_context_with_root to avoid calling the panicking task_scope_top() accessor when resolving the fallback parent UUID. Read the scope-stack top through the non-panicking lock path and convert a poisoned lock using the existing scope_stack_lock_error FlowError helper, preserving the declared Result-based API.crates/core/src/observability/mod.rs (1)
83-86: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winDuplicated doc summary on
default_mark_exclude_names.Two independent summary paragraphs precede the function; rustdoc renders both. Keep the rationale sentence or the "Return the..." sentence, not both.
📝 Proposed doc cleanup
/// Default mark names excluded from tool projection because they are emitted /// at high volume and are better represented as exporter-native events. -/// Return the default mark names excluded from OpenTelemetry projections. pub fn default_mark_exclude_names() -> Vec<String> {🤖 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/src/observability/mod.rs` around lines 83 - 86, Remove one of the two summary sentences in the documentation for default_mark_exclude_names, retaining either the rationale about high-volume marks or the return-value description so rustdoc has a single summary paragraph.
🤖 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 `@ATTRIBUTIONS-Rust.md`:
- Around line 10-16: The attribution generator produces Markdown blocks without
required spacing and fence language annotations. Update the template or
generation logic in generate_attributions.sh so headings and fenced license
blocks have surrounding blank lines and fences specify an appropriate language,
then regenerate ATTRIBUTIONS-Rust.md rather than editing it directly.
In `@crates/cli/src/sessions/mod.rs`:
- Around line 791-793: Refactor SessionManager::apply_events and the related
start_tool path so tool_conditional_execution(...).await runs after releasing
self.inner, rather than while the global session-map mutex is held. Preserve
event ordering using the existing per-session synchronization, and apply the
same ownership change to the additional tool-start handling at the referenced
location.
In `@crates/cli/tests/coverage/shared/plugins_tests.rs`:
- Around line 1236-1241: Update the test’s seed configuration to use the valid
Rampart field target_path_patterns instead of target_paths, then edit that field
through set_struct_field with both paths and assert the merged component
configuration retains the edited list. Keep future_top_level as the unknown-key
preservation case.
In `@crates/core/src/api/llm.rs`:
- Around line 543-579: Refactor emit_optimization_marks_with_async and its
callers to remove the unused async sanitize callback and unreachable None
branch, retaining a single mark-cursor loop with an enqueue closure of the
stated shape. Consolidate the duplicated contributions, mark_emitted, and break
logic from enqueue_optimization_marks and the test-only
emit_optimization_marks_with into that shared loop, then update the production,
reserved, and test seams to provide only their respective enqueue behavior.
- Around line 786-818: Extract the shared request-processing logic from
emit_llm_start_with_subscribers and this transform closure into a helper that
accepts an optional codec and returns (input, annotation). Reuse it in both
start paths, preserving credential stripping, sanitization, request-change
handling, and turn projection while passing each path’s actual codec instead of
hardcoding LlmSanitizeRequestContext::default() or None.
In `@crates/core/src/api/runtime/subscriber_dispatcher.rs`:
- Around line 812-822: Update the DispatcherMessage::Barrier handling to avoid
blocking the dispatcher indefinitely on publications.recv(). Use a bounded
receive-timeout loop that periodically logs the stalled event name and lineage,
then abandons the barrier after the configured limit while still releasing
permit; preserve publication handling when messages arrive and ensure all
dispatcher events and flush_subscribers() can continue despite a wedged
middleware future.
- Around line 567-576: Update in_dispatcher_callback to check whether the async
task-local publication buffer is active, matching the existing
THREAD_PUBLICATION_BUFFER branch, rather than treating any installed context as
active. Reuse the async buffer’s PublicationBuffer::is_active check so
with_async_publication_context with publication set to None does not cause
flush_subscribers to return early.
In `@crates/core/src/api/tool.rs`:
- Line 828: Update the Rust examples in the referenced documentation to await
calls to the async helpers tool_request_intercepts and
tool_conditional_execution. Add .await at each call site while preserving the
surrounding example logic and current API usage.
In `@crates/core/src/plugin/dynamic/native.rs`:
- Around line 3123-3126: Update the synchronous legacy ABI callback wrappers
around call_event_sanitize_callback and the additionally listed callback ranges
to execute blocking FFI work through spawn_blocking or the established
native_runtime() blocking pool before returning the future. Preserve callback
arguments, result propagation, and existing async ABI behavior while ensuring
v1/v2 plugin calls do not run directly on Tokio runtime workers.
- Around line 2099-2143: The native_string_from_json failure branch in the
chunk-processing loop must report an error terminal instead of breaking into the
successful completion callback. Update that None path to invoke cb with a
non-null error message and end-of-stream flag false, perform the same cleanup
via callback_guard.finish(), and return; preserve the existing successful
terminal callback for fully encoded streams.
- Around line 3968-3988: Refactor the stream polling flow around
poll_relay_llm_stream and spawn_with_continuation_context so each chunk does not
create and join a new OS thread or invoke a per-call block_on. Keep the stream
state and poisoned-lock error behavior intact while introducing a reusable
worker, runtime-driven stream task, or channel-based mechanism that persists
across poll calls and delivers chunks efficiently.
- Around line 1613-1657: Reclaim the raw Arc created for next_ref in both
failure branches of the native async callback flow: when catch_unwind handles a
panic and when NemoRelayNativeAsyncCallbackState::try_from returns Err. Drop it
with the matching NativeAsyncNext type before returning, while preserving the
existing completion and invocation cleanup.
In `@crates/core/src/plugin/dynamic/worker.rs`:
- Around line 1127-1136: The EventSanitizeFn closure currently ignores its
incoming EventSanitizeFields, causing worker sanitizers to overwrite prior chain
edits. Update the callback around invoke_event_sanitize to forward the captured
fields in the invocation payload, preserving earlier sanitizer changes while
keeping the existing event and callback_name handling unchanged.
In `@crates/core/src/stream.rs`:
- Around line 417-422: Update the Poll::Ready(Err(error)) branch of the
finalization task handling to preserve and return the stored terminal_result
instead of replacing it with FlowError::Internal. Log the join failure
represented by error, then clear finalization while surfacing the original
collector/provider error to the caller.
In `@crates/core/tests/fixtures/native_plugin/src/lib.rs`:
- Around line 22-27: Split ASYNC_PENDING_ENTERED into two independently named
AtomicBool flags: one for the pending tool-request path and one for the
cancellation probe in raw_async_tool_execution_callback. Update each setter and
its corresponding exported getter, then change
native_v3_async_registration_supports_all_middleware_kinds to poll the
appropriate getter for each phase so neither probe consumes the other’s state.
In `@crates/core/tests/integration/middleware_tests.rs`:
- Around line 2493-2494: Remove the unnecessary #[tokio::test] attributes from
test_concurrent_register_deregister, test_concurrent_intercept_mutations,
test_concurrent_register_and_read, both duplicate-registration tests, and both
deregistration tests, replacing them with plain #[test] attributes. Keep each
test body unchanged, including test_deregister_removes_from_chain’s
flush_subscribers usage.
- Around line 1379-1463: Ensure both affected cancellation tests restore the
originating thread scope stack even when assertions or other code paths fail.
Around each set_thread_scope_stack(create_scope_stack()) call, install the
existing ThreadScopeStackRestore-style guard using the captured originating
stack, and remove reliance on restoration statements that run only after
assertions; preserve the tests’ existing assertions and cleanup behavior.
In `@crates/core/tests/integration/pipeline_tests.rs`:
- Around line 2089-2104: Fix the timing race in the release-thread handshake
around flush_subscribers: ensure flush_started is recorded before the helper
thread begins its 50 ms delay. Update the synchronization between
release_started_tx/release_started_rx and the spawned release_thread so the
sleep starts strictly after flush_started, while preserving the assertion that
flush_subscribers waits for the pending stream END.
In `@crates/core/tests/integration/subscriber_dispatcher_tests.rs`:
- Around line 113-167: Add a concise comment immediately before the blocking
release_rx.lock().unwrap().recv() call in the sanitizer closure, stating that
the synchronous wait is intentional to block the dispatcher worker and that the
main test thread releases it, preserving the queued-event ordering assertion;
apply the same clarification to the analogous blocking section in the referenced
test range.
In `@crates/core/tests/unit/llm_api_tests.rs`:
- Around line 647-651: Move sanitizer input capture from the callback body into
the returned async future for the affected callbacks, including the cases around
the visible response callback and the additionally noted occurrences. Match the
poll-time recording pattern used by
failed_managed_calls_sanitize_fallback_end_data, while preserving each
callback’s existing result and assertions.
In `@crates/core/tests/unit/native_plugin_tests.rs`:
- Around line 1682-1750: Document the intentional two-worker runtime requirement
in the tests using invoke_native_next_then_return_state, noting that the
callback blocks while waiting for the continuation and therefore requires
another worker to poll it. Add the comment to both applicable test setups so
future changes do not reduce worker_threads(2) without understanding the
synchronization dependency.
In `@crates/core/tests/unit/runtime_state_tests.rs`:
- Around line 32-93: Extend each panic-containment test for the tool request,
tool response, LLM request, and LLM response sanitize chains with a second
well-behaved entry after the panicking entry. Assert the final result reflects
the later sanitizer’s output, proving each corresponding snapshot-chain method
continues processing entries after a panic while preserving the last valid
snapshot.
In `@crates/core/tests/unit/subscriber_dispatcher_tests.rs`:
- Around line 135-141: Increase the 100 ms recv_timeout in the flush timing
assertion to a generous duration that accommodates loaded CI runners, while
keeping the existing ordering: verify the flush result before releasing the
blocking delivery via release_tx and calling flush_subscribers.
In `@crates/ffi/src/callable.rs`:
- Around line 340-345: Update the callback wrappers that decode and free both
input and result pointers—especially the shown tool sanitizer flow, plus
wrap_tool_request_intercept_fn, wrap_tool_exec_fn, wrap_llm_exec_fn, and
wrap_llm_stream_exec_fn—to guard pointer identity before freeing either
allocation, matching wrap_llm_sanitize_response_fn. Preserve aliasing callbacks
by decoding the result safely and freeing shared pointers only once; apply the
same behavior consistently across all listed wrappers.
- Around line 982-995: Update the event sanitizer closure around the callback
invocation to clear the thread-local error before calling cb, then decode its
returned pointer with json_result_from_ptr instead of ptr_to_json so
callback-provided errors and null results propagate correctly. Preserve the
existing result deserialization and string cleanup, and update the affected
unit-test assertion for the new null-callback message.
In `@crates/ffi/tests/unit/api/core_tests.rs`:
- Around line 915-962: Extend
synchronous_ffi_middleware_helper_works_inside_tokio_with_scope_local_visibility
with a registered tool-request intercept callback that panics, then invoke the
FFI helper through the current-thread runtime. Assert the returned status is
NemoRelayStatus::Internal and verify the panic is contained rather than
unwinding across the C boundary, while preserving the existing scope cleanup.
In `@crates/ffi/tests/unit/api/registry_tests.rs`:
- Around line 417-418: Resolve the event sanitizer error policy consistently:
inspect invalid_event_sanitize_cb and the sanitizer handling in callable.rs,
then make failed sanitization drop the affected data and metadata fields rather
than returning the original payload, unless passthrough is explicitly intended.
Update the assertions for invalid_callback_event to verify the chosen
fail-closed behavior; if retaining passthrough, document that policy in the
sanitizer registration APIs and nemo_relay.h.
In `@crates/ffi/tests/unit/callable_tests.rs`:
- Around line 473-485: Replace the shared contains("invalid") assertion in the
loop over invalid_json_cb and invalid_utf8_cb with callback-specific expected
text. Assert that invalid_json_cb produces the wrapper’s “returned invalid JSON”
error and invalid_utf8_cb produces “returned invalid UTF-8”, matching the
messages generated through wrap_llm_sanitize_response_fn and
json_result_from_ptr.
In `@crates/node/pii_rampart.d.ts`:
- Around line 30-31: Update ConfigWithSelectors so both target_paths and
target_path_patterns require non-empty string tuples, such as [string,
...string[]], matching defaultConfig() validation and preventing empty selector
arrays at compile time.
In `@crates/node/pii_rampart.js`:
- Around line 49-58: Rename the exported Node helper ComponentSpec to
componentSpec, preserving its existing behavior and parameters. Update the
corresponding declaration in pii_rampart.d.ts and all references in
pii_rampart_tests.mjs so the public API consistently uses camelCase.
- Around line 29-45: Update the configuration assembly in defaultConfig so the
explicit modelPath argument remains authoritative: either apply config before
assigning model_path or exclude/reject config.model_path overrides. Preserve
other configuration overrides such as target_paths, while ensuring the returned
model_path always matches modelPath.
In `@crates/node/src/callable.rs`:
- Around line 220-232: Remove the unused private function
await_middleware_json_or_value from callable.rs. Keep
await_middleware_json_result and all active middleware handling unchanged.
- Around line 616-620: Update the queue-failure branch in the callback flow to
call record_callback_error(...) with the same failure message before returning
FlowError::Internal. Ensure getLastCallbackError() reflects failures to queue
the JS tool callback, matching the sanitize-wrapper behavior.
In `@crates/node/src/callback_factory.rs`:
- Around line 62-132: Update callPromise and its settlePublication lifecycle so
a callback that returns a never-settling promise cannot retain an active
publicationStates entry indefinitely; add an appropriate bounded cleanup
mechanism for the owned publication context while preserving normal
resolve/reject settlement behavior and avoiding cleanup of shared state.
In `@crates/node/tests/callback_error_tests.mjs`:
- Around line 145-156: Update the PromiseAwareFn conversion-failure test around
__testClosedPromiseAwareCall to call clearLastCallbackError() after the
rejection assertion, matching the sibling test’s cleanup and preventing the
recorded callback error from leaking into later getLastCallbackError()
assertions.
In `@crates/node/tests/event_sanitizers_tests.mjs`:
- Around line 346-352: Replace the wall-clock Promise.race assertion in the
queued-managed-tool execution test with a direct await of execution, then assert
inlineSanitizerReturned is false. Preserve the test’s verification that
execution completes while the sanitizer queue remains blocked, without relying
on a timing threshold.
In `@crates/node/tests/pii_rampart_tests.mjs`:
- Around line 21-24: Add coverage in the existing component construction test
around rampart.ComponentSpec and validate the disabled option by constructing
with { enabled: false }, then assert the resulting component’s enabled value is
false. Preserve the existing default-enabled assertions and configuration
diagnostics check.
In `@crates/node/tests/scope_local_tests.mjs`:
- Around line 783-803: Move the scope creation and both
scopeRegisterToolExecutionIntercept calls into the existing try block, making
pushScope the try block’s first operation. Keep the existing popScope cleanup in
finally so registration failures still remove the scope.
In `@crates/node/tests/test_support.mjs`:
- Around line 9-24: Update waitForSubscriberCallbacks to accept a descriptive
label for the wait and include that label in the timeout error; update its
anonymous-predicate call sites to provide identifying labels so CI logs reveal
which wait failed.
In `@crates/node/tests/tools_tests.mjs`:
- Around line 754-768: In the test flow around the tool call and event
assertions, replace the single flushSubscribers() wait with the imported
waitForSubscriberCallbacks helper so queued JS subscriber callbacks complete
before reading events. Keep the existing cleanup and assertions unchanged.
In `@crates/pii-redaction/README.md`:
- Around line 259-269: Update the selector behavior documentation in the README
after the normalized-shape explanation to state that, with the OpenAI Chat
codec, selectors touching choice-derived fields such as message, tool_calls,
finish_reason, or api_specific fail closed and omit the entire response when
multiple choices are returned.
In `@crates/pii-redaction/src/rampart/mod.rs`:
- Around line 603-617: Prevent selector configurations from accepting the empty
JSON pointer, which currently activates redaction without selecting object or
array content. Preserve is_valid_json_pointer’s empty-pointer behavior for other
callers, and add a non-empty validation check in config_value_violations for
every target_paths and target_path_patterns entry before
enforce_activation_invariants accepts the configuration.
In `@crates/pii-redaction/src/rampart/model.rs`:
- Around line 647-656: Replace the non-waiting try_inference_guard admission in
RampartSanitizer::sanitize_texts with bounded async waiting through a tokio
Semaphore, acquiring an owned permit before spawn_blocking and holding it
through inference. Add a timeout so concurrent requests queue briefly while
blocking threads remain bounded. Preserve payloads when admission succeeds, and
reserve the existing fail-closed full-redaction behavior for timeout exhaustion
only; update try_inference_guard or its callers accordingly.
In `@crates/plugin/README.md`:
- Around line 45-51: Update the migration-guide hyperlink in the “Stable native
ABI v3” documentation to use an anchor that actually exists on the referenced
page, replacing the invalid `#upgrade-to-nemo-relay-07` fragment while preserving
the link’s destination and surrounding text.
In `@crates/plugin/src/lib.rs`:
- Around line 2536-2541: Extract the shared v3 capability check and unsafe cast
into an `NemoRelayNativeHostApiV1::as_v3` helper returning
`Option<&NemoRelayNativeHostApiV3>`. Replace the duplicated validation and cast
in the raw registration paths and `OwnedHostApi::copy_from` with this accessor,
preserving `InvalidArg` handling when it returns `None`.
In `@crates/plugin/tests/typed_callbacks.rs`:
- Around line 36-66: The existing async ABI test only validates enum
discriminants and lacks coverage for the v3 capability gate. Extend the tests
around register_async_middleware_raw and register_async_stream_middleware_raw
using a v2 host-table variant derived from test_host() with abi_version 2 and
struct_size set to size_of::<NemoRelayNativeHostApiV1>(), asserting both
registrations return InvalidArg without reading beyond the table; also cover
successful registration and duplicate-name behavior with parity across both
registration paths.
In `@crates/python/src/py_api/mod.rs`:
- Around line 76-92: Add a regression test for block_on_sync_middleware using a
Tokio runtime configured with exactly one worker thread, and execute it from
within that runtime while the bridged future requires scheduler progress. Verify
the call completes successfully without deadlocking or starving the worker,
while preserving the existing current-thread coverage.
In `@crates/python/src/py_callable.rs`:
- Around line 62-94: Replace the poisoning-sensitive `expect("scheduled
awaitable")` mutex locks in both `CancellablePyFuture::poll` and
`CancellablePyFuture::drop` with recovery using `PoisonError::into_inner`,
preserving the existing task-clearing, cancellation, and completion behavior.
- Around line 770-833: Extract the repeated invocation bridge from
wrap_py_tool_fn, wrap_py_tool_conditional_fn, wrap_py_tool_request_intercept_fn,
wrap_py_tool_exec_fn, and the equivalent LLM wrappers into one shared helper.
Have it accept the callback, argument tuple, publication flag, invocation
context, and task locals, perform the existing context/publication dispatch, and
return the locals-aware split result. Replace each duplicated preamble and
four-arm dispatch with this helper while preserving current error handling and
invocation behavior.
In `@crates/python/src/py_types/observability.rs`:
- Around line 379-389: Update the force_flush documentation to explicitly state
that it must not be called from subscriber or middleware callbacks, and describe
that re-entrant calls do not establish the queued-delivery barrier. Keep the
existing behavior and sink-flush description unchanged.
In `@crates/python/tests/coverage/coverage_tests.rs`:
- Around line 664-757: Rename the enclosing test function
test_sync_wrapper_fallbacks_and_helpers to reflect that it now verifies wrapper
error propagation, while preserving all existing assertions and test behavior.
In `@crates/python/tests/coverage/nemo_guardrails_coverage_tests.rs`:
- Around line 318-334: Update the teardown drain flow around the embedded drain
module and run_until_complete call so drain failures are captured and handled
after the test result, allowing any original test panic to remain the reported
failure instead of being replaced by an unwrap panic. Also replace the escaped
one-line Python source passed to PyModule::from_code with a multiline raw string
literal while preserving the existing drain behavior.
In `@crates/python/tests/coverage/py_api_coverage_tests.rs`:
- Around line 928-937: Update
synchronous_middleware_bridge_avoids_tokio_runtime_reentry to execute its async
block on the pyo3_async_runtimes::tokio::get_runtime() handle, ensuring the test
exercises reentry into the runtime used by block_on_sync_middleware. Preserve
the existing result assertion; alternatively, rename the test to reflect only
that it offloads when running inside an arbitrary Tokio runtime.
- Around line 335-346: Update run_standalone so conditional_allowed reflects the
result of tool_conditional_execution rather than being assigned True
unconditionally; capture the call’s blocked/allowed outcome and preserve the
existing return field so the assertion meaningfully validates conditional
execution.
In `@crates/python/tests/coverage/py_callable_coverage_tests.rs`:
- Around line 21-90: Move InstalledContextModule and
install_event_sanitizer_context_module into tests/support alongside
init_python_test_locked, preserving their save/restore semantics. Update both
existing call sites to use this shared RAII helper and remove the duplicate
local implementation so the shared nemo_relay package setup is not shadowed.
- Around line 757-933: Add a cancellation test covering CancellablePyFuture:
start a Python awaitable that signals it has begun, drop the future before
completion, and assert the scheduled asyncio task finishes cancelled. Exercise
the unscheduled-awaitable cleanup path as needed, using the existing Python test
helpers and event-loop setup without altering current success or error tests.
In `@crates/python/tests/support/mod.rs`:
- Around line 27-32: Ensure every test in the crate, including
synchronous_middleware_bridge_avoids_tokio_runtime_reentry and
execution_next_context_restores_scope, acquires lock_python_test() before
executing. Keep the existing environment restoration and guard/drop ordering
unchanged so all process-global environment mutations remain serialized.
- Around line 65-69: The Python test guard’s isolated config directory is
configured but not lifecycle-managed. Update PythonTestGuard initialization and
its Drop implementation to create the temp directory before setting
XDG_CONFIG_HOME, then remove it when the last guard is dropped, preserving
shared-guard behavior and avoiding cleanup while other guards remain active.
In `@crates/worker/src/lib.rs`:
- Around line 135-136: Change the name parameter in ToolConditionalFn and
ToolRequestFn from owned String to borrowed &str, and apply the same signature
adjustment to the corresponding LlmRequestFn and related async handler
definitions around the additionally affected symbols. Update their call sites
and implementations to pass borrowed names while preserving ownership only where
required beyond the async invocation.
In `@docs/about-nemo-relay/concepts/subscribers.mdx`:
- Around line 162-165: Replace the one-event-loop-turn guidance with
predicate-based waiting for observable JavaScript callback delivery. Update
docs/about-nemo-relay/concepts/subscribers.mdx lines 162-165,
docs/getting-started/quick-start/nodejs.mdx lines 89-90, and
docs/instrument-applications/adding-scopes-and-marks.mdx lines 104-105 to wait
for the expected callback before validation, deregistration, or teardown; update
lines 190-191 to state that one event-loop turn is insufficient for validation.
In `@docs/instrument-applications/instrument-llm-call.mdx`:
- Around line 237-238: The async-flush guidance is incorrectly presented as
universal and interrupts the validation list. In
docs/instrument-applications/instrument-llm-call.mdx lines 237-238, scope
flushSubscribers() and the event-loop turn to Node.js, and add Python guidance
to await nemo_relay.subscribers.flush_async(). In
docs/instrument-applications/instrument-tool-call.mdx lines 197-198, make the
same Node.js/Python distinction for the example near line 78 and move the
paragraph outside the validation bullet list.
In `@docs/instrument-applications/instrument-tool-call.mdx`:
- Around line 197-198: Update the subscriber flush guidance in the
instrument-tool-call documentation to distinguish Node.js-only
flushSubscribers() usage and add the Python instruction to await
nemo_relay.subscribers.flush_async() inside asyncio.run(main()). Move the
inserted paragraph outside the existing validation bullet list so the list
remains contiguous, with its lead-in and parallel items intact.
In `@go/nemo_relay/pii_rampart.go`:
- Around line 57-64: Update RampartPiiComponent so callers can control the
component’s Enabled state instead of always receiving true, while preserving the
existing default behavior; add the corresponding options or disabled-component
helper and implement the same public API in the pii_rampart binding to keep Go
surfaces consistent.
- Around line 36-55: Update NewRampartPiiConfig so it cannot produce an
activation-invalid configuration: require and initialize at least one selector,
or add a selector-bearing constructor and make the existing path explicitly
non-activatable. Apply the same public API and behavior change across all
language bindings, preserving validation and registration invariants.
---
Outside diff comments:
In `@crates/core/src/api/runtime/scope_stack.rs`:
- Around line 446-458: Update capture_propagation_context_with_root to avoid
calling the panicking task_scope_top() accessor when resolving the fallback
parent UUID. Read the scope-stack top through the non-panicking lock path and
convert a poisoned lock using the existing scope_stack_lock_error FlowError
helper, preserving the declared Result-based API.
In `@crates/core/src/observability/mod.rs`:
- Around line 83-86: Remove one of the two summary sentences in the
documentation for default_mark_exclude_names, retaining either the rationale
about high-volume marks or the return-value description so rustdoc has a single
summary paragraph.
🪄 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: 010b959d-b6c8-49a4-b8ba-641fe5d5c263
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (139)
.gitattributesATTRIBUTIONS-Rust.mdCargo.tomlcrates/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/diagnostics/probes.rscrates/cli/src/plugins/editor_model.rscrates/cli/src/plugins/mod.rscrates/cli/src/server/mod.rscrates/cli/src/sessions/mod.rscrates/cli/tests/coverage/shared/plugins_tests.rscrates/cli/tests/coverage/shared/probes_tests.rscrates/cli/tests/coverage/shared/server_tests.rscrates/core/Cargo.tomlcrates/core/src/api/llm.rscrates/core/src/api/runtime.rscrates/core/src/api/runtime/callbacks.rscrates/core/src/api/runtime/continuation_context.rscrates/core/src/api/runtime/scope_stack.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/subscriber.rscrates/core/src/api/tool.rscrates/core/src/context/registries.rscrates/core/src/observability/mod.rscrates/core/src/plugin/dynamic/native.rscrates/core/src/plugin/dynamic/worker.rscrates/core/src/registry.rscrates/core/src/stream.rscrates/core/tests/fixtures/native_plugin/src/lib.rscrates/core/tests/fixtures/worker_plugin/src/main.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/continuation_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/observability/atof_tests.rscrates/core/tests/unit/observability/mod_tests.rscrates/core/tests/unit/plugin_tests.rscrates/core/tests/unit/runtime_state_tests.rscrates/core/tests/unit/shared_tests.rscrates/core/tests/unit/subscriber_dispatcher_tests.rscrates/ffi/README.mdcrates/ffi/nemo_relay.hcrates/ffi/src/api/llm_registry.rscrates/ffi/src/api/mod.rscrates/ffi/src/api/plugin.rscrates/ffi/src/callable.rscrates/ffi/tests/integration/callable_extra_tests.rscrates/ffi/tests/integration/main.rscrates/ffi/tests/support/mod.rscrates/ffi/tests/unit/api/core_tests.rscrates/ffi/tests/unit/api/registry_tests.rscrates/ffi/tests/unit/callable_tests.rscrates/node/README.mdcrates/node/package.jsoncrates/node/pii_rampart.d.tscrates/node/pii_rampart.jscrates/node/plugin.d.tscrates/node/src/api/mod.rscrates/node/src/callable.rscrates/node/src/callback_factory.rscrates/node/src/promise_call.rscrates/node/src/types/mod.rscrates/node/tests/adaptive_tests.mjscrates/node/tests/callback_error_tests.mjscrates/node/tests/event_sanitizers_tests.mjscrates/node/tests/llm_tests.mjscrates/node/tests/pii_rampart_tests.mjscrates/node/tests/rust/promise_call_tests.rscrates/node/tests/scope_local_tests.mjscrates/node/tests/scope_tests.mjscrates/node/tests/test_support.mjscrates/node/tests/tools_tests.mjscrates/pii-redaction/Cargo.tomlcrates/pii-redaction/README.mdcrates/pii-redaction/src/builtin.rscrates/pii-redaction/src/lib.rscrates/pii-redaction/src/rampart/mod.rscrates/pii-redaction/src/rampart/model.rscrates/pii-redaction/src/rampart/prefilter.rscrates/pii-redaction/src/rampart/sanitizer.rscrates/pii-redaction/src/rampart/tokenizer.rscrates/pii-redaction/tests/unit/component_tests.rscrates/plugin/README.mdcrates/plugin/src/lib.rscrates/plugin/tests/typed_callbacks.rscrates/python/src/lib.rscrates/python/src/py_api/mod.rscrates/python/src/py_callable.rscrates/python/src/py_types/core.rscrates/python/src/py_types/observability.rscrates/python/src/test_support.rscrates/python/tests/coverage/coverage_tests.rscrates/python/tests/coverage/nemo_guardrails_coverage_tests.rscrates/python/tests/coverage/py_api_coverage_tests.rscrates/python/tests/coverage/py_callable_coverage_tests.rscrates/python/tests/coverage/py_types_coverage_tests.rscrates/python/tests/support/mod.rscrates/worker/README.mdcrates/worker/src/lib.rscrates/worker/tests/worker_sdk_tests.rsdocs/about-nemo-relay/concepts/middleware.mdxdocs/about-nemo-relay/concepts/subscribers.mdxdocs/build-plugins/dynamic-plugins/native-dynamic/about.mdxdocs/getting-started/quick-start/nodejs.mdxdocs/instrument-applications/adding-scopes-and-marks.mdxdocs/instrument-applications/advanced-guide.mdxdocs/instrument-applications/instrument-llm-call.mdxdocs/instrument-applications/instrument-tool-call.mdxdocs/reference/event-sanitizers.mdxdocs/reference/migration-guides.mdxgo/nemo_relay/README.mdgo/nemo_relay/event_sanitizers_test.gogo/nemo_relay/llm/llm_shorthand_test.gogo/nemo_relay/llm_test.gogo/nemo_relay/nemo_relay.gogo/nemo_relay/pii_rampart.gogo/nemo_relay/pii_rampart/pii_rampart.gogo/nemo_relay/pii_rampart/pii_rampart_test.gogo/nemo_relay/pii_rampart_test.go
| config: json!({ | ||
| "version": 1, | ||
| "model_path": "/models/rampart", | ||
| "target_paths": ["/message"], | ||
| "future_top_level": "preserve" | ||
| }) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Seed key target_paths is not a Rampart schema field — likely meant target_path_patterns.
Line 436 asserts the schema exposes target_path_patterns, so target_paths here is just a second unknown key alongside future_top_level. The test therefore never exercises merge_known_editor_object for the List field, which is the interesting merge path. Use the real field name and assert its edited value survives.
♻️ Cover the list-field merge path
config: json!({
"version": 1,
"model_path": "/models/rampart",
- "target_paths": ["/message"],
+ "target_path_patterns": ["/message"],
"future_top_level": "preserve"
})Then assert the edited list persists:
set_struct_field(
&mut rampart.config,
"target_path_patterns",
json!(["/message", "/metadata/note"]),
)
.unwrap();
// ...
assert_eq!(
component.config.get("target_path_patterns"),
Some(&json!(["/message", "/metadata/note"]))
);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| config: json!({ | |
| "version": 1, | |
| "model_path": "/models/rampart", | |
| "target_paths": ["/message"], | |
| "future_top_level": "preserve" | |
| }) | |
| config: json!({ | |
| "version": 1, | |
| "model_path": "/models/rampart", | |
| "target_path_patterns": ["/message"], | |
| "future_top_level": "preserve" | |
| }) |
🤖 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/cli/tests/coverage/shared/plugins_tests.rs` around lines 1236 - 1241,
Update the test’s seed configuration to use the valid Rampart field
target_path_patterns instead of target_paths, then edit that field through
set_struct_field with both paths and assert the merged component configuration
retains the edited list. Keep future_top_level as the unknown-key preservation
case.
Source: Path instructions
| export type ConfigWithSelectors = Partial<Config> & | ||
| ({ target_paths: string[] } | { target_path_patterns: string[] }); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Model selectors as non-empty arrays.
defaultConfig() rejects missing or empty selectors, but string[] allows [] and both union branches can be empty at compile time. Use a tuple such as [string, ...string[]] so TypeScript rejects configurations that will always throw at runtime.
🤖 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/node/pii_rampart.d.ts` around lines 30 - 31, Update
ConfigWithSelectors so both target_paths and target_path_patterns require
non-empty string tuples, such as [string, ...string[]], matching defaultConfig()
validation and preventing empty selector arrays at compile time.
| return { | ||
| version: 1, | ||
| model_path: modelPath, | ||
| input: true, | ||
| output: true, | ||
| mark: true, | ||
| tool_input: true, | ||
| tool_output: true, | ||
| priority: 100, | ||
| target_paths: [], | ||
| target_path_patterns: [], | ||
| min_score: 0.4, | ||
| excluded_labels: [], | ||
| replacement: '[REDACTED]', | ||
| max_windows_per_payload: 128, | ||
| inference_batch_size: 16, | ||
| ...config, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep the explicit modelPath authoritative.
Because ...config appears after model_path: modelPath, defaultConfig('/approved', { model_path: '/other', target_paths: ['/x'] }) returns /other. This silently loads a different model than the function argument indicates. Reject model_path in overrides or spread configuration before assigning the authoritative path; update the declaration accordingly.
🤖 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/node/pii_rampart.js` around lines 29 - 45, Update the configuration
assembly in defaultConfig so the explicit modelPath argument remains
authoritative: either apply config before assigning model_path or exclude/reject
config.model_path overrides. Preserve other configuration overrides such as
target_paths, while ensuring the returned model_path always matches modelPath.
| /** | ||
| * Wrap Rampart PII config as a top-level plugin component. | ||
| * | ||
| * @param {object} config - Rampart PII component configuration. | ||
| * @param {{ enabled?: boolean }} [options={}] - Optional component flags. | ||
| * @returns {object} A shared plugin component spec. | ||
| */ | ||
| function ComponentSpec(config, { enabled = true } = {}) { | ||
| return plugin.ComponentSpec(RAMPART_PII_PLUGIN_KIND, config, { enabled }); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Use camelCase for the exported Node helper.
ComponentSpec is a public Node function, but the binding convention requires camelCase public APIs. Rename it to componentSpec and update crates/node/pii_rampart.d.ts and crates/node/tests/pii_rampart_tests.mjs together.
As per path instructions, Node.js public APIs must use camelCase.
🤖 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/node/pii_rampart.js` around lines 49 - 58, Rename the exported Node
helper ComponentSpec to componentSpec, preserving its existing behavior and
parameters. Update the corresponding declaration in pii_rampart.d.ts and all
references in pii_rampart_tests.mjs so the public API consistently uses
camelCase.
Source: Path instructions
| const component = rampart.ComponentSpec(config); | ||
| assert.equal(component.kind, rampart.RAMPART_PII_PLUGIN_KIND); | ||
| assert.equal(component.enabled, true); | ||
| assert.deepEqual(rampart.validateConfig(config).diagnostics, []); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Add coverage for disabled component construction.
The test only verifies the default enabled: true path. Add a componentSpec(config, { enabled: false }) assertion so the public option and its serialized value are protected.
As per path instructions, tests should cover the behavior promised by the changed API surface.
🤖 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/node/tests/pii_rampart_tests.mjs` around lines 21 - 24, Add coverage
in the existing component construction test around rampart.ComponentSpec and
validate the disabled option by constructing with { enabled: false }, then
assert the resulting component’s enabled value is false. Preserve the existing
default-enabled assertions and configuration diagnostics check.
Source: Path instructions
| target_paths = ["/message"] | ||
| target_path_patterns = [ | ||
| "/messages/*/content", | ||
| "/messages/*/content/*/text", | ||
| ] | ||
| ``` | ||
|
|
||
| `target_paths` contains exact JSON pointers. `target_path_patterns` also | ||
| accepts `*` as one complete path segment. At least one selector is required. | ||
| When a supported codec is active, selectors address the normalized Relay | ||
| request or response shape. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the multi-choice omission behavior for choice-derived selectors.
/message is a choice-derived normalized field, so with codec = "openai_chat" this example config causes sanitize_response_with_codec to omit the entire response payload whenever a provider returns more than one choice (crates/pii-redaction/src/rampart/sanitizer.rs Lines 378-386). Readers copying this config will see responses disappear with no explanation. Add a sentence after Line 269 noting that selectors touching message, tool_calls, finish_reason, or api_specific fail closed for multi-choice OpenAI Chat responses.
🤖 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/pii-redaction/README.md` around lines 259 - 269, Update the selector
behavior documentation in the README after the normalized-shape explanation to
state that, with the OpenAI Chat codec, selectors touching choice-derived fields
such as message, tool_calls, finish_reason, or api_specific fail closed and omit
the entire response when multiple choices are returned.
| fn is_valid_json_pointer(path: &str) -> bool { | ||
| if path.is_empty() { | ||
| return true; | ||
| } | ||
| if !path.starts_with('/') { | ||
| return false; | ||
| } | ||
| let mut bytes = path.as_bytes().iter().copied(); | ||
| while let Some(byte) = bytes.next() { | ||
| if byte == b'~' && !matches!(bytes.next(), Some(b'0' | b'1')) { | ||
| return false; | ||
| } | ||
| } | ||
| true | ||
| } |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Empty pointer passes selector validation but selects nothing.
is_valid_json_pointer("") returns true, so target_paths = [""] satisfies the "must select explicit content fields" invariant at Line 491 and passes enforce_activation_invariants. compile_json_pointer then produces an empty segment vector, which only matches a root-level JSON string — never an object/array payload — so the component activates and redacts nothing. That is a silent fail-open for a config that looks valid.
🛡️ Reject the whole-document pointer for selectors
fn is_valid_json_pointer(path: &str) -> bool {
if path.is_empty() {
- return true;
+ return false;
}If the empty pointer must stay valid for other callers, add a dedicated non-empty check in config_value_violations for both target_paths and target_path_patterns instead.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fn is_valid_json_pointer(path: &str) -> bool { | |
| if path.is_empty() { | |
| return true; | |
| } | |
| if !path.starts_with('/') { | |
| return false; | |
| } | |
| let mut bytes = path.as_bytes().iter().copied(); | |
| while let Some(byte) = bytes.next() { | |
| if byte == b'~' && !matches!(bytes.next(), Some(b'0' | b'1')) { | |
| return false; | |
| } | |
| } | |
| true | |
| } | |
| fn is_valid_json_pointer(path: &str) -> bool { | |
| if path.is_empty() { | |
| return false; | |
| } | |
| if !path.starts_with('/') { | |
| return false; | |
| } | |
| let mut bytes = path.as_bytes().iter().copied(); | |
| while let Some(byte) = bytes.next() { | |
| if byte == b'~' && !matches!(bytes.next(), Some(b'0' | b'1')) { | |
| return false; | |
| } | |
| } | |
| true | |
| } |
🤖 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/pii-redaction/src/rampart/mod.rs` around lines 603 - 617, Prevent
selector configurations from accepting the empty JSON pointer, which currently
activates redaction without selecting object or array content. Preserve
is_valid_json_pointer’s empty-pointer behavior for other callers, and add a
non-empty validation check in config_value_violations for every target_paths and
target_path_patterns entry before enforce_activation_invariants accepts the
configuration.
| fn try_inference_guard(lock: &Mutex<()>) -> PluginResult<MutexGuard<'_, ()>> { | ||
| lock.try_lock().map_err(|error| match error { | ||
| TryLockError::WouldBlock => { | ||
| PluginError::Internal("Rampart inference is already running".into()) | ||
| } | ||
| TryLockError::Poisoned(error) => { | ||
| PluginError::Internal(format!("Rampart inference lock poisoned: {error}")) | ||
| } | ||
| }) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Non-waiting admission turns any concurrent request into a full redaction of every selected field.
try_inference_guard returns Internal on WouldBlock, and RampartSanitizer::sanitize_texts treats that error as fail-closed by replacing all selected texts with replacement. With one global mutex and zero waiting, two overlapping requests mean the loser's entire selected content is destroyed — not delayed. Under any real concurrency this silently mangles observability payloads, and it contradicts the PR's reported "queueing behind the single model mutex" behavior.
Since the callbacks are now async and inference runs in spawn_blocking, prefer bounded waiting over instant failure: hold a tokio::sync::Semaphore (or an owned permit acquired before entering spawn_blocking) with a timeout, and reserve fail-closed for the timeout case only. That keeps blocking threads bounded while preserving payload fidelity.
🤖 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/pii-redaction/src/rampart/model.rs` around lines 647 - 656, Replace
the non-waiting try_inference_guard admission in
RampartSanitizer::sanitize_texts with bounded async waiting through a tokio
Semaphore, acquiring an owned permit before spawn_blocking and holding it
through inference. Add a timeout so concurrent requests queue briefly while
blocking threads remain bounded. Preserve payloads when admission succeeds, and
reserve the existing fail-closed full-redaction behavior for timeout exhaustion
only; update try_inference_guard or its callers accordingly.
| // NewRampartPiiConfig returns Rampart PII settings with runtime defaults. | ||
| func NewRampartPiiConfig(modelPath string) RampartPiiConfig { | ||
| return RampartPiiConfig{ | ||
| Version: 1, | ||
| ModelPath: modelPath, | ||
| Input: true, | ||
| Output: true, | ||
| Mark: true, | ||
| ToolInput: true, | ||
| ToolOutput: true, | ||
| Priority: 100, | ||
| TargetPaths: []string{}, | ||
| TargetPathPatterns: []string{}, | ||
| MinScore: 0.4, | ||
| ExcludedLabels: []string{}, | ||
| Replacement: "[REDACTED]", | ||
| MaxWindowsPerPayload: 128, | ||
| InferenceBatchSize: 16, | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not return an activation-invalid default configuration.
NewRampartPiiConfig() returns both selector lists empty, while the Rampart contract requires at least one selector. Callers must mutate the result before validation or registration, unlike the Node helper, which enforces this invariant. Add a selector-bearing constructor or make the invalid intermediate state explicit.
As per path instructions, binding changes are public API changes and should preserve behavior across language bindings.
🤖 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 `@go/nemo_relay/pii_rampart.go` around lines 36 - 55, Update
NewRampartPiiConfig so it cannot produce an activation-invalid configuration:
require and initialize at least one selector, or add a selector-bearing
constructor and make the existing path explicitly non-activatable. Apply the
same public API and behavior change across all language bindings, preserving
validation and registration invariants.
Source: Path instructions
| // RampartPiiComponent converts config into the shared plugin component. | ||
| func RampartPiiComponent(config RampartPiiConfig) PluginComponentSpec { | ||
| return PluginComponentSpec{ | ||
| Kind: RampartPiiPluginKind, | ||
| Enabled: true, | ||
| Config: mustConfigMap(config), | ||
| } | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Expose the component enabled state consistently across bindings.
The Node API supports ComponentSpec(config, { enabled }), but the Go helper hardcodes Enabled: true. Go callers cannot construct a disabled component or round-trip the shared component shape. Add an options or disabled-component helper and mirror it in go/nemo_relay/pii_rampart/pii_rampart.go.
As per path instructions, binding changes must preserve parity across public language surfaces.
🤖 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 `@go/nemo_relay/pii_rampart.go` around lines 57 - 64, Update
RampartPiiComponent so callers can control the component’s Enabled state instead
of always receiving true, while preserving the existing default behavior; add
the corresponding options or disabled-component helper and implement the same
public API in the pii_rampart binding to keep Go surfaces consistent.
Source: Path instructions
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
🛑 Comments failed to post (54)
ATTRIBUTIONS-Rust.md (1)
10-16: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix Markdown formatting in the attribution generator.
The regenerated blocks repeatedly trigger MD022, MD031, and MD040: headings and fenced license blocks lack surrounding blank lines, and fences omit a language. Update
./scripts/generate_attributions.shor its template, then regenerate this file; do not hand-edit the generated artifact.🧰 Tools
🪛 markdownlint-cli2 (0.23.1)
[warning] 12-12: Headings should be surrounded by blank lines
Expected: 1; Actual: 0; Below(MD022, blanks-around-headings)
[warning] 15-15: Headings should be surrounded by blank lines
Expected: 1; Actual: 0; Above(MD022, blanks-around-headings)
[warning] 15-15: Headings should be surrounded by blank lines
Expected: 1; Actual: 0; Below(MD022, blanks-around-headings)
[warning] 16-16: Fenced code blocks should be surrounded by blank lines
(MD031, blanks-around-fences)
[warning] 16-16: Fenced code blocks should have a language specified
(MD040, fenced-code-language)
🤖 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 `@ATTRIBUTIONS-Rust.md` around lines 10 - 16, The attribution generator produces Markdown blocks without required spacing and fence language annotations. Update the template or generation logic in generate_attributions.sh so headings and fenced license blocks have surrounding blank lines and fences specify an appropriate language, then regenerate ATTRIBUTIONS-Rust.md rather than editing it directly.Source: Linters/SAST tools
crates/cli/src/sessions/mod.rs (1)
791-793: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Release the manager lock before awaiting conditional guardrails.
tool_conditional_execution(...).awaitexecutes plugin-defined async work whileSessionManager::apply_eventsretainsself.innerthrough event application. A slow guardrail blocks all sessions and can deadlock reentrant hook handling. Refactor ownership so this await occurs outside the global session-map mutex, preserving ordering with per-session synchronization.Also applies to: 1518-1536
🤖 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/cli/src/sessions/mod.rs` around lines 791 - 793, Refactor SessionManager::apply_events and the related start_tool path so tool_conditional_execution(...).await runs after releasing self.inner, rather than while the global session-map mutex is held. Preserve event ordering using the existing per-session synchronization, and apply the same ownership change to the additional tool-start handling at the referenced location.crates/core/tests/fixtures/native_plugin/src/lib.rs (1)
22-27: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Single shared flag serves two distinct probes.
ASYNC_PENDING_ENTEREDis set both by the pending tool-request path (Line 928) and by the cancellation probe inraw_async_tool_execution_callback(Line 1046), and the exported getter resets it. Becausenative_v3_async_registration_supports_all_middleware_kindspolls the same flag for both phases, a stale set from the first phase can satisfy the second wait. The initialassert!(!pending_entered())plus the reset makes the current sequence work, but two named flags would make the coupling explicit and the test non-order-dependent.Also applies to: 927-958
🤖 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/fixtures/native_plugin/src/lib.rs` around lines 22 - 27, Split ASYNC_PENDING_ENTERED into two independently named AtomicBool flags: one for the pending tool-request path and one for the cancellation probe in raw_async_tool_execution_callback. Update each setter and its corresponding exported getter, then change native_v3_async_registration_supports_all_middleware_kinds to poll the appropriate getter for each phase so neither probe consumes the other’s state.crates/core/tests/integration/middleware_tests.rs (2)
1379-1463: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Restore the thread scope stack even if the assertions fail.
set_thread_scope_stack(create_scope_stack())at Lines 1441/1579 replaces the ambient stack, and the original is only restored after theassert_eq!on the sanitizedEndevent. On assertion failure the thread keeps the throwaway stack, and the still-heldTEST_MUTEXis poisoned, so subsequent tests in this binary see a confusing cascade instead of the real failure. A scope-guard (asnative_plugin_tests.rsdoes withThreadScopeStackRestore) would keep the failure localized.Also applies to: 1520-1604
🤖 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/integration/middleware_tests.rs` around lines 1379 - 1463, Ensure both affected cancellation tests restore the originating thread scope stack even when assertions or other code paths fail. Around each set_thread_scope_stack(create_scope_stack()) call, install the existing ThreadScopeStackRestore-style guard using the captured originating stack, and remove reliance on restoration statements that run only after assertions; preserve the tests’ existing assertions and cleanup behavior.
2493-2494: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
These tests gained
#[tokio::test]without awaiting anything.
test_concurrent_register_deregister,test_concurrent_intercept_mutations,test_concurrent_register_and_read, the two duplicate-registration tests, and the two deregistration tests contain no.await, so the runtime is dead weight (andtest_deregister_removes_from_chainonly awaits nothing either — it usesflush_subscribers). Keeping them as plain#[test]avoids implying async semantics where registration is purely synchronous.Also applies to: 2539-2540, 2583-2584, 3334-3335, 3365-3366, 3402-3403, 3415-3416
🤖 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/integration/middleware_tests.rs` around lines 2493 - 2494, Remove the unnecessary #[tokio::test] attributes from test_concurrent_register_deregister, test_concurrent_intercept_mutations, test_concurrent_register_and_read, both duplicate-registration tests, and both deregistration tests, replacing them with plain #[test] attributes. Keep each test body unchanged, including test_deregister_removes_from_chain’s flush_subscribers usage.crates/core/tests/integration/pipeline_tests.rs (1)
2089-2104: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Race in the timing assertion:
elapsed()can legitimately be under 50 ms.The helper thread sends
()and then starts its 50 ms sleep, so the release fires atsend + 50ms.flush_startedis recorded afterrecv_timeoutreturns, i.e. atsend + deltawithdelta > 0. The flush therefore unblocks after only50ms - delta, and the>= 50msassertion fails whenever the main thread's wakeup is delayed. Invert the handshake so the sleep starts strictly afterflush_startedis taken.🐛 Proposed fix: let the releaser wait for the main thread
- let (release_started_tx, release_started_rx) = std::sync::mpsc::channel(); - let release_thread = std::thread::spawn(move || { - release_started_tx.send(()).unwrap(); - std::thread::sleep(std::time::Duration::from_millis(50)); - sanitizer_release.notify_one(); - }); - release_started_rx - .recv_timeout(std::time::Duration::from_secs(2)) - .expect("sanitizer release thread did not start"); - let flush_started = std::time::Instant::now(); + let (go_tx, go_rx) = std::sync::mpsc::channel(); + let release_thread = std::thread::spawn(move || { + go_rx + .recv_timeout(std::time::Duration::from_secs(2)) + .expect("main thread did not start flushing"); + std::thread::sleep(std::time::Duration::from_millis(50)); + sanitizer_release.notify_one(); + }); + let flush_started = std::time::Instant::now(); + go_tx.send(()).unwrap(); flush_subscribers().unwrap();📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.let (go_tx, go_rx) = std::sync::mpsc::channel(); let release_thread = std::thread::spawn(move || { go_rx .recv_timeout(std::time::Duration::from_secs(2)) .expect("main thread did not start flushing"); std::thread::sleep(std::time::Duration::from_millis(50)); sanitizer_release.notify_one(); }); let flush_started = std::time::Instant::now(); go_tx.send(()).unwrap(); flush_subscribers().unwrap(); assert!( flush_started.elapsed() >= std::time::Duration::from_millis(50), "flush must wait for the pending stream END" ); release_thread.join().unwrap();🤖 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/integration/pipeline_tests.rs` around lines 2089 - 2104, Fix the timing race in the release-thread handshake around flush_subscribers: ensure flush_started is recorded before the helper thread begins its 50 ms delay. Update the synchronization between release_started_tx/release_started_rx and the spawned release_thread so the sleep starts strictly after flush_started, while preserving the assertion that flush_subscribers waits for the pending stream END.crates/core/tests/integration/subscriber_dispatcher_tests.rs (1)
113-167: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Deliberate blocking inside the async sanitizer — worth a one-line comment.
release_rx.lock().unwrap().recv()blocks the dispatcher worker rather than awaiting, which is precisely what makes the "already-queued event" ordering observable. Since that inverts the usual async guidance, a short comment stating that the block is intentional (and that only the main thread can release it) will stop a future reader from "fixing" it into an async channel and silently weakening the test.Also applies to: 169-240
🤖 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/integration/subscriber_dispatcher_tests.rs` around lines 113 - 167, Add a concise comment immediately before the blocking release_rx.lock().unwrap().recv() call in the sanitizer closure, stating that the synchronous wait is intentional to block the dispatcher worker and that the main test thread releases it, preserving the queued-event ordering assertion; apply the same clarification to the analogous blocking section in the referenced test range.crates/core/tests/unit/llm_api_tests.rs (1)
647-651: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Note: capture happens eagerly, outside the returned future.
sanitizer_inputs.lock().unwrap().push(response)runs when the callback is invoked, not when the future is polled. That is fine for these assertions, but it no longer exercises the awaited path — contrast with thefailed_managed_calls_sanitize_fallback_end_dataversion at Lines 1579-1586, which records inside the async block. Consider aligning the two so at least one buffered/streaming case verifies capture at poll time.Also applies to: 676-681, 777-782, 808-814
🤖 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/llm_api_tests.rs` around lines 647 - 651, Move sanitizer input capture from the callback body into the returned async future for the affected callbacks, including the cases around the visible response callback and the additionally noted occurrences. Match the poll-time recording pattern used by failed_managed_calls_sanitize_fallback_end_data, while preserving each callback’s existing result and assertions.crates/core/tests/unit/native_plugin_tests.rs (1)
1682-1750: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Verify the 2-worker requirement is intentional and sufficient.
invoke_native_next_then_return_stateblocks onrecv_timeoutinside the callback, so the invoked continuation must be polled by a different worker. Withworker_threads(2)this holds, but it is an implicit coupling: if the callback ever runs on the same worker that would poll the continuation, the test degrades to a 1s timeout and then assertsinvoke_status == Okon a continuation that never started. A short comment stating why two workers are required would prevent a future flake being misdiagnosed.Also applies to: 1751-1828
🤖 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/native_plugin_tests.rs` around lines 1682 - 1750, Document the intentional two-worker runtime requirement in the tests using invoke_native_next_then_return_state, noting that the callback blocks while waiting for the continuation and therefore requires another worker to poll it. Add the comment to both applicable test setups so future changes do not reduce worker_threads(2) without understanding the synchronization dependency.crates/core/tests/unit/runtime_state_tests.rs (1)
32-93: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Single-entry chains don't pin the continuation guarantee.
Each sanitize chain is exercised with exactly one panicking entry, so the tests confirm "panic is contained" but not the documented "preserve the last valid snapshot and keep running the remaining entries". Adding a second, well-behaved entry after the panicking one (as
dispatcher_publishes_the_snapshot_when_an_async_sanitizer_panicsdoes for marks) would catch a regression that aborts the chain on panic.♻️ Example for the tool request sanitize chain
let tool_entries = vec![RegistryRecord::new("tool-panic", 0, tool_sanitizer)]; + let tool_followup: ToolSanitizeFn = Arc::new(|_, mut args| { + Box::pin(async move { + args["followup"] = json!(true); + Ok(args) + }) + }); + let tool_entries = { + let mut entries = tool_entries; + entries.push(RegistryRecord::new("tool-followup", 1, tool_followup)); + entries + }; assert_eq!( NemoRelayContextState::tool_sanitize_request_snapshot_chain( "tool", tool_payload.clone(), &tool_entries, ) .await, - tool_payload + json!({"tool": "preserved", "followup": true}) );📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.let tool_payload = json!({"tool": "preserved"}); let tool_sanitizer: ToolSanitizeFn = Arc::new(|_, _| Box::pin(async { panic!("tool sanitizer panic") })); let tool_entries = vec![RegistryRecord::new("tool-panic", 0, tool_sanitizer)]; let tool_followup: ToolSanitizeFn = Arc::new(|_, mut args| { Box::pin(async move { args["followup"] = json!(true); Ok(args) }) }); let tool_entries = { let mut entries = tool_entries; entries.push(RegistryRecord::new("tool-followup", 1, tool_followup)); entries }; assert_eq!( NemoRelayContextState::tool_sanitize_request_snapshot_chain( "tool", tool_payload.clone(), &tool_entries, ) .await, json!({"tool": "preserved", "followup": true}) ); let tool_response = json!({"tool_response": "preserved"}); let tool_response_sanitizer: ToolSanitizeFn = Arc::new(|_, _| Box::pin(async { panic!("tool response sanitizer panic") })); assert_eq!( NemoRelayContextState::tool_sanitize_response_snapshot_chain( "tool", tool_response.clone(), &[ RegistryRecord::new( "tool-response-panic", 0, tool_response_sanitizer, ), ], ) .await, tool_response ); let request = LlmRequest { headers: Map::new(), content: json!({"llm": "preserved"}), }; let llm_sanitizer: LlmSanitizeRequestFn = Arc::new(|_, _| Box::pin(async { panic!("LLM sanitizer panic") })); let llm_entries = vec![RegistryRecord::new("llm-panic", 0, llm_sanitizer)]; assert_eq!( NemoRelayContextState::llm_sanitize_request_snapshot_chain( request.clone(), LlmSanitizeRequestContext::default(), &llm_entries, ) .await, Some(request.clone()) ); let llm_response = json!({"llm_response": "preserved"}); let llm_response_sanitizer: LlmSanitizeResponseFn = Arc::new(|_, _| Box::pin(async { panic!("LLM response sanitizer panic") })); assert_eq!( NemoRelayContextState::llm_sanitize_response_snapshot_chain( llm_response.clone(), LlmSanitizeResponseContext::default(), &[ RegistryRecord::new( "llm-response-panic", 0, llm_response_sanitizer, ), ], ) .await, Some(llm_response) );🤖 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/runtime_state_tests.rs` around lines 32 - 93, Extend each panic-containment test for the tool request, tool response, LLM request, and LLM response sanitize chains with a second well-behaved entry after the panicking entry. Assert the final result reflects the later sanitizer’s output, proving each corresponding snapshot-chain method continues processing entries after a panic while preserving the last valid snapshot.crates/core/tests/unit/subscriber_dispatcher_tests.rs (1)
135-141: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Tighten-bound flake risk: this is a positive timing assertion.
Unlike the negative assertion in the previous test (where the flush provably cannot complete), here the flush must be observed within 100 ms. On a loaded CI runner the dispatcher may simply not be scheduled in time, failing a test that is semantically correct. The bound can be generous without weakening the assertion, since the blocking delivery is only released afterwards at Line 136.
🩹 Proposed fix
- let flush_result = flush_rx.recv_timeout(std::time::Duration::from_millis(100)); + let flush_result = flush_rx.recv_timeout(std::time::Duration::from_secs(5));📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.let flush_result = flush_rx.recv_timeout(std::time::Duration::from_secs(5)); let _ = release_tx.send(()); flush_subscribers().unwrap(); assert!( flush_result.is_ok(), "a delivery queued after a flush must not delay that flush" );🤖 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/subscriber_dispatcher_tests.rs` around lines 135 - 141, Increase the 100 ms recv_timeout in the flush timing assertion to a generous duration that accommodates loaded CI runners, while keeping the existing ordering: verify the flush result before releasing the blocking delivery via release_tx and calling flush_subscribers.crates/ffi/src/callable.rs (2)
340-345: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Input string is freed before the returned pointer is decoded; no aliasing guard here, unlike the response sanitizer.
c_argsis freed at Line 342, thenresult_ptris decoded at Line 343 and freed at Line 344. If a C callback returns the pointer it was handed, this is a read-after-free plus a double-free.wrap_llm_sanitize_response_fnexplicitly tolerates that case (if result_ptr != response_jsonat Line 813) andcrates/ffi/tests/unit/callable_tests.rsexercises alias callbacks against it, so aliasing is a pattern this crate already treats as reachable — just not uniformly.Either make the ABI contract explicit ("the returned string must be a fresh allocation") in the doc comment plus
crates/ffi/nemo_relay.h, or add the same identity guard. The identical shape appears inwrap_tool_request_intercept_fn(Lines 391-396),wrap_tool_exec_fn(Lines 413-417),wrap_llm_exec_fn(Lines 881-885), andwrap_llm_stream_exec_fn(Lines 907-911).🛡️ Guard variant, if aliasing should stay permitted
let result_ptr = unsafe { cb(ud.ptr, c_name.as_ptr(), c_args) }; - unsafe { nemo_relay_string_free_internal(c_args) }; let result = json_result_from_ptr(result_ptr, "tool sanitize callback returned null"); - unsafe { nemo_relay_string_free_internal(result_ptr) }; + unsafe { + nemo_relay_string_free_internal(c_args); + if result_ptr != c_args { + nemo_relay_string_free_internal(result_ptr); + } + } result📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.let c_args = json_to_c_string(&args); let result_ptr = unsafe { cb(ud.ptr, c_name.as_ptr(), c_args) }; let result = json_result_from_ptr(result_ptr, "tool sanitize callback returned null"); unsafe { nemo_relay_string_free_internal(c_args); if result_ptr != c_args { nemo_relay_string_free_internal(result_ptr); } } result🤖 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/ffi/src/callable.rs` around lines 340 - 345, Update the callback wrappers that decode and free both input and result pointers—especially the shown tool sanitizer flow, plus wrap_tool_request_intercept_fn, wrap_tool_exec_fn, wrap_llm_exec_fn, and wrap_llm_stream_exec_fn—to guard pointer identity before freeing either allocation, matching wrap_llm_sanitize_response_fn. Preserve aliasing callbacks by decoding the result safely and freeing shared pointers only once; apply the same behavior consistently across all listed wrappers.
982-995: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Event sanitizer path discards the callback's own error message.
Unlike every other wrapper touched here, this one skips
clear_last_error()and decodes withptr_to_jsoninstead ofjson_result_from_ptr. A callback that sets an error and returns null therefore surfaces the generic "invalid event sanitizer result" rather than its own message, and a staleLAST_ERRORfrom a prior call on the same thread is never cleared.♻️ Align with the other wrappers
Box::pin(async move { + clear_last_error(); let ffi_event = FfiEvent((*event).clone()); let fields_json = json_to_c_string(&serde_json::to_value(&fields).unwrap_or(Json::Null)); let result_ptr = unsafe { cb(ud.ptr, &ffi_event, fields_json) }; unsafe { nemo_relay_string_free_internal(fields_json) }; - let result = serde_json::from_value(ptr_to_json(result_ptr)).map_err(|error| { - FlowError::Internal(format!("invalid event sanitizer result: {error}")) - }); + let result = json_result_from_ptr(result_ptr, "invalid event sanitizer result") + .and_then(|value| { + serde_json::from_value(value).map_err(|error| { + FlowError::Internal(format!("invalid event sanitizer result: {error}")) + }) + }); unsafe { nemo_relay_string_free_internal(result_ptr) }; result })Note this changes the null-callback message, so
crates/ffi/tests/unit/callable_tests.rslines 694-700 would need its assertion updated.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.Arc::new(move |event: Arc<Event>, fields: EventSanitizeFields| { let ud = ud.clone(); Box::pin(async move { clear_last_error(); let ffi_event = FfiEvent((*event).clone()); let fields_json = json_to_c_string(&serde_json::to_value(&fields).unwrap_or(Json::Null)); let result_ptr = unsafe { cb(ud.ptr, &ffi_event, fields_json) }; unsafe { nemo_relay_string_free_internal(fields_json) }; let result = json_result_from_ptr(result_ptr, "invalid event sanitizer result") .and_then(|value| { serde_json::from_value(value).map_err(|error| { FlowError::Internal(format!("invalid event sanitizer result: {error}")) }) }); unsafe { nemo_relay_string_free_internal(result_ptr) }; result })🤖 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/ffi/src/callable.rs` around lines 982 - 995, Update the event sanitizer closure around the callback invocation to clear the thread-local error before calling cb, then decode its returned pointer with json_result_from_ptr instead of ptr_to_json so callback-provided errors and null results propagate correctly. Preserve the existing result deserialization and string cleanup, and update the affected unit-test assertion for the new null-callback message.crates/ffi/tests/unit/api/core_tests.rs (1)
915-962: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Good targeted test; the panic-mapping branch of
block_on_sync_ffistays uncovered.This correctly pins the scope-local visibility that
with_scope_stackprovides across the helper-thread hop, and using a current-thread runtime exercises theHandle::try_current().is_ok()path. TheFlowError::Internal("...helper thread panicked...")arm atcrates/ffi/src/api/mod.rslines 120-125 has no test, though — a registered intercept callback that panics would assert both that the status maps toInternaland that the panic does not unwind across the C boundary.Want me to draft that case?
🤖 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/ffi/tests/unit/api/core_tests.rs` around lines 915 - 962, Extend synchronous_ffi_middleware_helper_works_inside_tokio_with_scope_local_visibility with a registered tool-request intercept callback that panics, then invoke the FFI helper through the current-thread runtime. Assert the returned status is NemoRelayStatus::Internal and verify the panic is contained rather than unwinding across the C boundary, while preserving the existing scope cleanup.crates/ffi/tests/unit/api/registry_tests.rs (1)
417-418: 🔒 Security & Privacy | 🟠 Major | 🏗️ Heavy lift
This encodes a fail-open change for event sanitizers — confirm that is intended.
invalid_event_sanitize_cb(Line 107) returns"not-json". Previously the test asserteddataandmetadatawereJson::Null, i.e. a malfunctioning sanitizer caused the fields to be dropped. The new expectation is the untouched original payload{"secret": true}from Lines 238-239, so a broken sanitizer now lets the raw payload reach subscribers.That follows from
crates/ffi/src/callable.rsLines 990-992 returningErrwhere the wrapper previously yielded default fields. For a redaction-oriented surface this inverts the failure posture, and it sits awkwardly beside the fail-closed guarantees stated for the Rampart plugin in this PR. If passthrough-on-error is the deliberate policy, it should be documented on the sanitizer registration APIs and incrates/ffi/nemo_relay.h; if not, a failed event sanitizer should drop the fields it was supposed to sanitize.#!/bin/bash # How does core handle an Err from an event sanitizer chain? rg -nP -C6 '(mark_sanitize|scope_sanitize|event_sanitize).*chain|sanitize_event' crates/core/src/api/shared.rs crates/core/src/api/scope.rs # Look for documented failure-posture wording on sanitizer errors. rg -nP -C3 'sanitiz\w+' docs/reference/event-sanitizers.mdx🤖 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/ffi/tests/unit/api/registry_tests.rs` around lines 417 - 418, Resolve the event sanitizer error policy consistently: inspect invalid_event_sanitize_cb and the sanitizer handling in callable.rs, then make failed sanitization drop the affected data and metadata fields rather than returning the original payload, unless passthrough is explicitly intended. Update the assertions for invalid_callback_event to verify the chosen fail-closed behavior; if retaining passthrough, document that policy in the sanitizer registration APIs and nemo_relay.h.crates/ffi/tests/unit/callable_tests.rs (1)
473-485: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Assertion is too loose to distinguish the two failure modes.
Both
invalid_json_cbandinvalid_utf8_cbare checked againstcontains("invalid"), so a regression collapsing UTF-8 and JSON failures into one message — or misattributing either — still passes.json_result_from_ptremits distinguishable"{fallback}: invalid UTF-8"and"{fallback}: invalid JSON"text; assert against it.💚 Tighter per-callback assertions
- for callback in [invalid_json_cb, invalid_utf8_cb] { + for (callback, expected) in [ + (invalid_json_cb, "invalid JSON"), + (invalid_utf8_cb, "invalid UTF-8"), + ] { let malformed_response = wrap_llm_sanitize_response_fn(callback, std::ptr::null_mut(), None); let error = resolve(malformed_response( json!({"secret": "must be preserved"}), nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), )) .unwrap_err(); assert!( - error.to_string().contains("invalid"), + error.to_string().contains(expected), "unexpected sanitizer error: {error}" ); }Note
wrap_llm_sanitize_response_fnproduces "returned invalid UTF-8"/"returned invalid JSON" (Lines 799-810), so adjust the expected strings to whichever wrapper each callback exercises.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.for (callback, expected) in [ (invalid_json_cb, "invalid JSON"), (invalid_utf8_cb, "invalid UTF-8"), ] { let malformed_response = wrap_llm_sanitize_response_fn(callback, std::ptr::null_mut(), None); let error = resolve(malformed_response( json!({"secret": "must be preserved"}), nemo_relay::api::runtime::LlmSanitizeResponseContext::default(), )) .unwrap_err(); assert!( error.to_string().contains(expected), "unexpected sanitizer error: {error}" ); }🤖 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/ffi/tests/unit/callable_tests.rs` around lines 473 - 485, Replace the shared contains("invalid") assertion in the loop over invalid_json_cb and invalid_utf8_cb with callback-specific expected text. Assert that invalid_json_cb produces the wrapper’s “returned invalid JSON” error and invalid_utf8_cb produces “returned invalid UTF-8”, matching the messages generated through wrap_llm_sanitize_response_fn and json_result_from_ptr.crates/node/src/callable.rs (2)
220-232: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash rg -nP --type=rust '\bawait_middleware_json_or_value\b'Repository: NVIDIA/NeMo-Relay
Length of output: 155
🏁 Script executed:
#!/bin/bash set -euo pipefail # Inspect the surrounding implementation and all in-file references. sed -n '1,320p' crates/node/src/callable.rs | nl -ba | sed -n '180,260p' printf '\n---\n' rg -n '\bawait_middleware_json_(or_value|result)\b' crates/node/src/callable.rs printf '\n---\n' rg -n '\bawait_middleware_json_or_value\b' crates/node/srcRepository: NVIDIA/NeMo-Relay
Length of output: 195
🏁 Script executed:
#!/bin/bash set -euo pipefail # Show the relevant section with portable line numbering. awk 'NR>=180 && NR<=260 { printf "%4d %s\n", NR, $0 }' crates/node/src/callable.rs printf '\n---\n' # Find all local references to the helper and its sibling. rg -n '\bawait_middleware_json_(or_value|result)\b' crates/node/src/callable.rs crates/node/srcRepository: NVIDIA/NeMo-Relay
Length of output: 4978
Remove the unused fallback helper
crates/node/src/callable.rs:220-232—await_middleware_json_or_valuehas no callers now, so the private function will tripdead_codeundercargo clippy --workspace --all-targets -- -D warnings.🤖 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/node/src/callable.rs` around lines 220 - 232, Remove the unused private function await_middleware_json_or_value from callable.rs. Keep await_middleware_json_result and all active middleware handling unchanged.Source: Coding guidelines
616-620: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash rg -nP -C3 '__testClosedToolCallback' crates/node/src rg -nP -C2 'record_callback_error' crates/node/src/callable.rs | head -60Repository: NVIDIA/NeMo-Relay
Length of output: 2872
🏁 Script executed:
#!/bin/bash set -euo pipefail # Show the surrounding call sites in crates/node/src/callable.rs sed -n '580,700p' crates/node/src/callable.rs printf '\n---\n' sed -n '800,880p' crates/node/src/callable.rs printf '\n---\n' sed -n '1030,1075p' crates/node/src/callable.rs printf '\n=== callback error plumbing ===\n' sed -n '1,220p' crates/node/src/convert.rs printf '\n=== relevant test ===\n' sed -n '1,160p' crates/node/tests/callback_error_tests.mjsRepository: NVIDIA/NeMo-Relay
Length of output: 17571
🏁 Script executed:
#!/bin/bash set -euo pipefail rg -n "record_callback_error|getLastCallbackError|__testClosedToolCallback|failed to queue JS tool callback|queue JS .*callback" crates/node/src crates/node/tests -n -C 2Repository: NVIDIA/NeMo-Relay
Length of output: 34449
Record queue failures here too
These branches returnFlowErrorwithout callingrecord_callback_error(...), unlike the sanitize wrappers. That leavesgetLastCallbackError()stale when a callback can’t be queued; record the message before returning.🤖 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/node/src/callable.rs` around lines 616 - 620, Update the queue-failure branch in the callback flow to call record_callback_error(...) with the same failure message before returning FlowError::Internal. Ensure getLastCallbackError() reflects failures to queue the JS tool callback, matching the sanitize-wrapper behavior.crates/node/src/callback_factory.rs (1)
62-132: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value
publicationStatesentries leak if a callback promise never settles.
settlePublicationis only reached from thethen/catchhandlers, so a JS middleware callback that returns a forever-pending promise permanently retains its map entry (and keepsactive === true, which suppresses flush waiting for that context id). Given callbacks are user-supplied, consider bounding this or documenting the contract.🤖 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/node/src/callback_factory.rs` around lines 62 - 132, Update callPromise and its settlePublication lifecycle so a callback that returns a never-settling promise cannot retain an active publicationStates entry indefinitely; add an appropriate bounded cleanup mechanism for the owned publication context while preserving normal resolve/reject settlement behavior and avoiding cleanup of shared state.crates/node/tests/callback_error_tests.mjs (1)
145-156: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Clear the recorded callback error to avoid cross-test leakage.
The sibling test at line 83 calls
clearLastCallbackError(). If the forced conversion failure also records into the callback-error slot, a latergetLastCallbackError()assertion in this suite can observe it.💚 Reset the error slot
assert.equal(invoked, false); + clearLastCallbackError(); });📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.it('PromiseAwareFn argument conversion failures reject without invoking the callback', async () => { let invoked = false; await assert.rejects( () => __testClosedPromiseAwareCall(() => { invoked = true; return null; }, true), /forced PromiseAwareFn conversion failure/i, ); assert.equal(invoked, false); clearLastCallbackError(); });🤖 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/node/tests/callback_error_tests.mjs` around lines 145 - 156, Update the PromiseAwareFn conversion-failure test around __testClosedPromiseAwareCall to call clearLastCallbackError() after the rejection assertion, matching the sibling test’s cleanup and preventing the recorded callback error from leaking into later getLastCallbackError() assertions.Source: Path instructions
crates/node/tests/event_sanitizers_tests.mjs (1)
346-352: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Timing-based assertion will flake on loaded CI.
Promise.raceagainst a 250 ms timer asserts a positive (tool execution finishes fast) rather than a negative, so a slow runner turns this into a spurious failure. Prefer proving non-blocking behavior without a wall-clock budget, e.g. await the execution and assertinlineSanitizerReturned === falseat that point (the blocker still holds the sanitizer queue), which is what the test actually cares about.♻️ Suggested deterministic form
- const execution = lib.toolCallExecute('queued-managed-tool', {}, (args) => args); - const executionState = await Promise.race([ - execution.then(() => 'executed'), - new Promise((resolve) => setTimeout(() => resolve('blocked'), 250)), - ]); - assert.equal(executionState, 'executed'); - assert.equal(inlineSanitizerReturned, false); + await lib.toolCallExecute('queued-managed-tool', {}, (args) => args); + assert.equal(inlineSanitizerReturned, false);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.await lib.toolCallExecute('queued-managed-tool', {}, (args) => args); assert.equal(inlineSanitizerReturned, false);🤖 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/node/tests/event_sanitizers_tests.mjs` around lines 346 - 352, Replace the wall-clock Promise.race assertion in the queued-managed-tool execution test with a direct await of execution, then assert inlineSanitizerReturned is false. Preserve the test’s verification that execution completes while the sanitizer queue remains blocked, without relying on a timing threshold.crates/node/tests/scope_local_tests.mjs (1)
783-803: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Scope push and registrations sit outside the
try.If either
scopeRegisterToolExecutionInterceptcall throws,popScopenever runs and the leaked scope pollutes the shared scope stack for subsequent tests in this file. Move line 783 onward inside thetry(or push the scope as the first statement of thetry).🤖 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/node/tests/scope_local_tests.mjs` around lines 783 - 803, Move the scope creation and both scopeRegisterToolExecutionIntercept calls into the existing try block, making pushScope the try block’s first operation. Keep the existing popScope cleanup in finally so registration failures still remove the scope.crates/node/tests/test_support.mjs (1)
9-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Add a label to the timeout error.
Ten-plus call sites pass anonymous predicates; a bare "timed out waiting for subscriber callbacks" gives no hint which wait failed in CI logs.
♻️ Suggested change
-export async function waitForSubscriberCallbacks(predicate, timeoutMs = 15000) { +export async function waitForSubscriberCallbacks(predicate, timeoutMs = 15000, label = '') { await flushSubscribers(); @@ if (Date.now() >= deadline) { - throw new Error('timed out waiting for subscriber callbacks'); + throw new Error(`timed out waiting for subscriber callbacks${label ? `: ${label}` : ''}`); }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.export async function waitForSubscriberCallbacks(predicate, timeoutMs = 15000, label = '') { await flushSubscribers(); // flushSubscribers() waits for Relay's Rust subscriber dispatcher, but JS // subscriber callbacks are queued onto Node's event loop through N-API // ThreadsafeFunction. Yield event-loop turns until the observed JS-side // callback state is ready, with a timeout to avoid hanging the test forever. const deadline = Date.now() + timeoutMs; while (!predicate()) { await flushSubscribers(); if (Date.now() >= deadline) { throw new Error(`timed out waiting for subscriber callbacks${label ? `: ${label}` : ''}`); } await new Promise((resolve) => setImmediate(resolve)); } await flushSubscribers(); }🤖 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/node/tests/test_support.mjs` around lines 9 - 24, Update waitForSubscriberCallbacks to accept a descriptive label for the wait and include that label in the timeout error; update its anonymous-predicate call sites to provide identifying labels so CI logs reveal which wait failed.crates/node/tests/tools_tests.mjs (1)
754-768: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Use the new helper here; a single
flushSubscribers()doesn't guarantee the JS subscriber ran.
flushSubscribers()drains the Rust dispatcher, but the JS subscriber callback is queued onto Node's event loop through the threadsafe function — exactly the gapwaitForSubscriberCallbacks(imported at line 7) was added to close. Readingeventsat lines 765-768 right after one flush can observe an empty array and fail with aTypeErroronstart.data.♻️ Suggested change
const handle = toolCall('node_manual_tool_flush', { original: true }); toolCallEnd(handle, { ok: true }); - await flushSubscribers(); + await waitForSubscriberCallbacks( + () => + events.some((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'start') && + events.some((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'end'), + );📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.try { const handle = toolCall('node_manual_tool_flush', { original: true }); toolCallEnd(handle, { ok: true }); await waitForSubscriberCallbacks( () => events.some((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'start') && events.some((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'end'), ); } finally { deregisterToolSanitizeRequestGuardrail('node_manual_tool_flush_request'); deregisterToolSanitizeResponseGuardrail('node_manual_tool_flush_response'); deregisterSubscriber('node_manual_tool_flush_subscriber'); } assert.equal(requestFlushed, true); assert.equal(responseFlushed, true); const start = events.find((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'start'); const end = events.find((event) => event.name === 'node_manual_tool_flush' && event.scope_category === 'end'); assert.deepEqual(start.data, { original: true, requestSanitized: true }); assert.deepEqual(end.data, { ok: true, responseSanitized: true });🤖 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/node/tests/tools_tests.mjs` around lines 754 - 768, In the test flow around the tool call and event assertions, replace the single flushSubscribers() wait with the imported waitForSubscriberCallbacks helper so queued JS subscriber callbacks complete before reading events. Keep the existing cleanup and assertions unchanged.crates/plugin/README.md (1)
45-51: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash curl -sSIL -o /dev/null -w '%{http_code} %{url_effective}\n' \ "https://docs.nvidia.com/nemo/relay/reference/migration-guides"Repository: NVIDIA/NeMo-Relay
Length of output: 221
🏁 Script executed:
python3 - <<'PY' import urllib.request, re url = "https://docs.nvidia.com/nemo/relay/reference/migration-guides" html = urllib.request.urlopen(url, timeout=20).read().decode("utf-8", "replace") patterns = [ r'id="upgrade-to-nemo-relay-07"', r'name="upgrade-to-nemo-relay-07"', r'href="`#upgrade-to-nemo-relay-07`"', r'upgrade-to-nemo-relay-07', ] for p in patterns: m = re.search(p, html) print(f"{p}: {'YES' if m else 'NO'}") PYRepository: NVIDIA/NeMo-Relay
Length of output: 2441
🏁 Script executed:
python3 - <<'PY' import urllib.request, re, sys url = "https://docs.nvidia.com/nemo/relay/reference/migration-guides" req = urllib.request.Request(url, headers={"User-Agent":"Mozilla/5.0"}) html = urllib.request.urlopen(req, timeout=20).read().decode("utf-8", "replace") for needle in ["upgrade-to-nemo-relay-07", "Upgrade To NeMo Relay 0.7", "Migration Guides"]: print(needle, "YES" if needle in html else "NO") m = re.search(r'id="([^"]*upgrade-to-nemo-relay-07[^"]*)"', html) print("ID_MATCH:", m.group(1) if m else "NONE") PYRepository: NVIDIA/NeMo-Relay
Length of output: 2441
🌐 Web query:
site:docs.nvidia.com/nemo/relay/reference/migration-guides "upgrade-to-nemo-relay-07"💡 Result:
There is no official documentation or public information regarding an upgrade to a version designated as nemo-relay-07. The existing official NVIDIA NeMo Relay migration documentation specifically outlines the process for upgrading to version 0.6 from version 0.5 [1]. To manage upgrades for NeMo Relay, you should consult the official NVIDIA NeMo Relay Migration Guides, which provide sequential steps if you are skipping multiple versions [1]. General administrative tasks for the NeMo Relay CLI, such as validating or managing plugins, are performed using the nemo-relay command line interface [2]. If you are looking to update your installation, ensure you are referencing the specific version requirements and migration paths defined in the official documentation for the version you intend to reach [1].
Citations:
- 1: https://docs.nvidia.com/nemo/relay/reference/migration-guides
- 2: https://docs.nvidia.com/nemo/relay/dev/configure-plugins/discoverable-plugins
🏁 Script executed:
python3 - <<'PY' import subprocess, re, sys, tempfile, os url = "https://docs.nvidia.com/nemo/relay/reference/migration-guides" cmd = ["curl", "-ksSL", url] html = subprocess.check_output(cmd, text=True, timeout=30) hits = [] for pat in [r'id="upgrade-to-nemo-relay-07"', r'name="upgrade-to-nemo-relay-07"', r'href="`#upgrade-to-nemo-relay-07`"', r'upgrade-to-nemo-relay-07']: hits.append((pat, bool(re.search(pat, html)))) print("\n".join(f"{p}: {'YES' if ok else 'NO'}" for p, ok in hits)) # show nearby context if found m = re.search(r'.{0,120}upgrade-to-nemo-relay-07.{0,120}', html, re.I | re.S) if m: print("\nCONTEXT:\n", m.group(0)) PYRepository: NVIDIA/NeMo-Relay
Length of output: 291
Fix the
#upgrade-to-nemo-relay-07fragment. The migration guide page does not expose that anchor, so this link lands nowhere.🤖 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/README.md` around lines 45 - 51, Update the migration-guide hyperlink in the “Stable native ABI v3” documentation to use an anchor that actually exists on the referenced page, replacing the invalid `#upgrade-to-nemo-relay-07` fragment while preserving the link’s destination and surrounding text.Source: Coding guidelines
crates/plugin/src/lib.rs (1)
2536-2541: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Extract the v3 capability check into one helper.
The
abi_version/struct_sizegate plus theNemoRelayNativeHostApiV3cast is now duplicated in both raw registrations and again inOwnedHostApi::copy_from(Line 3260). One accessor keeps the unsafe cast and its precondition in a single place.♻️ Suggested helper
impl NemoRelayNativeHostApiV1 { fn as_v3(&self) -> Option<&NemoRelayNativeHostApiV3> { (self.abi_version >= NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE && self.struct_size >= std::mem::size_of::<NemoRelayNativeHostApiV3>()) .then(|| unsafe { &*(self as *const _ as *const NemoRelayNativeHostApiV3) }) } }- if self.host.abi_version < NEMO_RELAY_NATIVE_ABI_VERSION_ASYNC_MIDDLEWARE - || self.host.struct_size < std::mem::size_of::<NemoRelayNativeHostApiV3>() - { - return NemoRelayStatus::InvalidArg; - } - let host = unsafe { &*(self.host as *const _ as *const NemoRelayNativeHostApiV3) }; + let Some(host) = self.host.as_v3() else { + return NemoRelayStatus::InvalidArg; + };Also applies to: 2572-2577
🤖 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/src/lib.rs` around lines 2536 - 2541, Extract the shared v3 capability check and unsafe cast into an `NemoRelayNativeHostApiV1::as_v3` helper returning `Option<&NemoRelayNativeHostApiV3>`. Replace the duplicated validation and cast in the raw registration paths and `OwnedHostApi::copy_from` with this accessor, preserving `InvalidArg` handling when it returns `None`.crates/plugin/tests/typed_callbacks.rs (1)
36-66: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
Add coverage for the v3 capability gate.
Discriminant round-tripping is covered, but nothing exercises
register_async_middleware_raw/register_async_stream_middleware_rawagainst a legacy host table (abi_version = 2,struct_size = size_of::<NemoRelayNativeHostApiV1>()), which is the path that must returnInvalidArginstead of reading past the table.test_host()already exists in this file, so a v2 variant is cheap.As per coding guidelines, "Add tests covering registration and duplicate names, ... and parity across affected bindings."
🤖 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 36 - 66, The existing async ABI test only validates enum discriminants and lacks coverage for the v3 capability gate. Extend the tests around register_async_middleware_raw and register_async_stream_middleware_raw using a v2 host-table variant derived from test_host() with abi_version 2 and struct_size set to size_of::<NemoRelayNativeHostApiV1>(), asserting both registrations return InvalidArg without reading beyond the table; also cover successful registration and duplicate-name behavior with parity across both registration paths.Source: Coding guidelines
crates/worker/src/lib.rs (1)
135-136: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Owned
Stringname is inconsistent with the other async surfaces.
ToolConditionalFn/ToolRequestFn/LlmRequestFnnow take an owned name, while the already-async sanitize and execution surfaces keep&str(Lines 425, 447, 518, 631). Line 1745 shows a borrowed name works fine with a future-returning handler, since the payload outlives theawaitin the same function, so this is an avoidable, gratuitous breaking change in the public SDK signature. Prefer&strfor all name parameters unless something forces ownership.Also applies to: 264-273
🤖 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/worker/src/lib.rs` around lines 135 - 136, Change the name parameter in ToolConditionalFn and ToolRequestFn from owned String to borrowed &str, and apply the same signature adjustment to the corresponding LlmRequestFn and related async handler definitions around the additionally affected symbols. Update their call sites and implementations to pass borrowed names while preserving ownership only where required beyond the async invocation.docs/about-nemo-relay/concepts/subscribers.mdx (1)
162-165: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Do Not Treat One Event-Loop Turn as a Delivery Barrier
flushSubscribers()completes Rust-side dispatch, but JavaScript callbacks can still be queued through N-API. The Node test helper polls until the callback state is observable, so onesetImmediatecan leave examples and assertions flaky. Use a condition-based wait before validation or deregistration.
docs/about-nemo-relay/concepts/subscribers.mdx#L162-L165: describe a predicate-based wait rather than one tick.docs/getting-started/quick-start/nodejs.mdx#L89-L90: wait for the expected callback before deregistering.docs/instrument-applications/adding-scopes-and-marks.mdx#L104-L105: wait for observed delivery before teardown.docs/instrument-applications/adding-scopes-and-marks.mdx#L190-L191: do not present one event-loop turn as sufficient for validation.As per path instructions, documentation must remain technically accurate against the current API.
📍 Affects 3 files
docs/about-nemo-relay/concepts/subscribers.mdx#L162-L165(this comment)docs/getting-started/quick-start/nodejs.mdx#L89-L90docs/instrument-applications/adding-scopes-and-marks.mdx#L104-L105docs/instrument-applications/adding-scopes-and-marks.mdx#L190-L191🤖 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 `@docs/about-nemo-relay/concepts/subscribers.mdx` around lines 162 - 165, Replace the one-event-loop-turn guidance with predicate-based waiting for observable JavaScript callback delivery. Update docs/about-nemo-relay/concepts/subscribers.mdx lines 162-165, docs/getting-started/quick-start/nodejs.mdx lines 89-90, and docs/instrument-applications/adding-scopes-and-marks.mdx lines 104-105 to wait for the expected callback before validation, deregistration, or teardown; update lines 190-191 to state that one event-loop turn is insufficient for validation.Source: Path instructions
docs/instrument-applications/instrument-llm-call.mdx (1)
237-238: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Copy-pasted async-flush note covers only the Node.js binding. Both pages gained the same paragraph telling readers to await
flushSubscribers(), a Node.js-only symbol, while the Python examples on those pages still call the synchronousnemo_relay.subscribers.flush()from insideasyncio.run(main())— whichdocs/reference/event-sanitizers.mdx(lines 81–84) anddocs/reference/migration-guides.mdx(lines 119–123) document as raising.
docs/instrument-applications/instrument-llm-call.mdx#L237-L238: scope the sentence to Node.js and add the Pythonawait nemo_relay.subscribers.flush_async()guidance matching the example at line 102.docs/instrument-applications/instrument-tool-call.mdx#L197-L198: apply the same scoping and Python guidance for the example at line 78, and move the paragraph out of the middle of the validation bullet list.📍 Affects 2 files
docs/instrument-applications/instrument-llm-call.mdx#L237-L238(this comment)docs/instrument-applications/instrument-tool-call.mdx#L197-L198🤖 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 `@docs/instrument-applications/instrument-llm-call.mdx` around lines 237 - 238, The async-flush guidance is incorrectly presented as universal and interrupts the validation list. In docs/instrument-applications/instrument-llm-call.mdx lines 237-238, scope flushSubscribers() and the event-loop turn to Node.js, and add Python guidance to await nemo_relay.subscribers.flush_async(). In docs/instrument-applications/instrument-tool-call.mdx lines 197-198, make the same Node.js/Python distinction for the example near line 78 and move the paragraph outside the validation bullet list.docs/instrument-applications/instrument-tool-call.mdx (1)
197-198: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Same Node-only flush guidance, and here it splits the validation bullet list.
Two problems in this insertion:
flushSubscribers()is Node.js-only, and the Python example on this page (line 78) callsnemo_relay.subscribers.flush()from insideasyncio.run(main()), whichdocs/reference/migration-guides.mdx(lines 119–123) documents as raising. Add theawait nemo_relay.subscribers.flush_async()instruction and scope the Node wording.- The paragraph is inserted between bullet items (line 195 and line 199), breaking one list into two fragments. Move it above or below the list.
As per coding guidelines, "All lists must have a complete lead-in sentence, more than one item, no more than two levels, parallel sentence construction".
🤖 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 `@docs/instrument-applications/instrument-tool-call.mdx` around lines 197 - 198, Update the subscriber flush guidance in the instrument-tool-call documentation to distinguish Node.js-only flushSubscribers() usage and add the Python instruction to await nemo_relay.subscribers.flush_async() inside asyncio.run(main()). Move the inserted paragraph outside the existing validation bullet list so the list remains contiguous, with its lead-in and parallel items intact.Source: Coding guidelines
crates/core/src/api/llm.rs (2)
543-579: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
The async
sanitizehook is dead abstraction, and the mark-cursor loop now exists in three copies.Both production callers pass
|event| async { Some(event) }, so theF/Futgeneric and theNone => break"sanitizer context unavailable" branch are unreachable; the two callers differ only in theirenqueueclosure. Meanwhileenqueue_optimization_marks(586-601) and the#[cfg(test)]emit_optimization_marks_with(677-700) reimplement the same contributions/mark_emitted/breakcursor logic, so a future fix to the accounting rule has to land in three places.Suggest dropping the
sanitizeparameter, keeping oneenqueue: impl FnMut(&Event, &[EventSubscriberFn]) -> boolloop, and having the sync and test seams call that single loop.Also applies to: 603-611
🤖 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/src/api/llm.rs` around lines 543 - 579, Refactor emit_optimization_marks_with_async and its callers to remove the unused async sanitize callback and unreachable None branch, retaining a single mark-cursor loop with an enqueue closure of the stated shape. Consolidate the duplicated contributions, mark_emitted, and break logic from enqueue_optimization_marks and the test-only emit_optimization_marks_with into that shared loop, then update the production, reserved, and test seams to provide only their respective enqueue behavior.
786-818: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Credential stripping, request sanitization, and turn projection are now duplicated across the two start paths.
This transform closure re-implements
emit_llm_start_with_subscribers(lines 436-476) with subtly different codec handling (LlmSanitizeRequestContext::default()andNonecodec for projection). Since one of the duplicated steps is the security-relevant credential removal documented at lines 731-739, the two copies must not drift. Extract a shared helper that takes an optional codec and returns(input, annotation).🤖 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/src/api/llm.rs` around lines 786 - 818, Extract the shared request-processing logic from emit_llm_start_with_subscribers and this transform closure into a helper that accepts an optional codec and returns (input, annotation). Reuse it in both start paths, preserving credential stripping, sanitization, request-change handling, and turn projection while passing each path’s actual codec instead of hardcoding LlmSanitizeRequestContext::default() or None.crates/core/src/api/runtime/subscriber_dispatcher.rs (2)
567-576: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Inconsistent activity check makes
flush_subscribersa silent no-op.The thread-local branch filters on
PublicationBuffer::is_active, but the task-local branch only tests for the presence of the task-local.with_async_publication_contextinstalls an inactive buffer whenpublicationisNone(line 622), so any code inside that scope seesin_dispatcher_callback() == trueandflush_subscribers()returnsOk(())without flushing, even though nothing is buffering its publications.🔒️ Align both branches on `is_active`
pub(super) fn in_dispatcher_callback() -> bool { IN_DISPATCHER.with(Cell::get) - || ASYNC_PUBLICATION_BUFFER.try_with(|_| ()).is_ok() + || ASYNC_PUBLICATION_BUFFER + .try_with(PublicationBuffer::is_active) + .unwrap_or(false) || THREAD_PUBLICATION_BUFFER.with(|buffer| { buffer .borrow() .as_ref() .is_some_and(PublicationBuffer::is_active) }) }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.pub(super) fn in_dispatcher_callback() -> bool { IN_DISPATCHER.with(Cell::get) || ASYNC_PUBLICATION_BUFFER .try_with(PublicationBuffer::is_active) .unwrap_or(false) || THREAD_PUBLICATION_BUFFER.with(|buffer| { buffer .borrow() .as_ref() .is_some_and(PublicationBuffer::is_active) }) }🤖 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/src/api/runtime/subscriber_dispatcher.rs` around lines 567 - 576, Update in_dispatcher_callback to check whether the async task-local publication buffer is active, matching the existing THREAD_PUBLICATION_BUFFER branch, rather than treating any installed context as active. Reuse the async buffer’s PublicationBuffer::is_active check so with_async_publication_context with publication set to None does not cause flush_subscribers to return early.
812-822: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Barrierblocks the single dispatcher thread on an unboundedrecv().
publications.recv()has no timeout. Sender drop covers panics and dropped futures, but a middleware future that is polled and never settles (a JS promise that never resolves, a worker plugin that never replies) holds theAsyncPublicationsender alive indefinitely. In that state the dispatcher stops delivering all events and everyflush_subscribers()blocks forever, with no diagnostic.Consider a bounded
recv_timeoutloop that logs a stall (event name + lineage) and eventually abandons the barrier, so one wedged middleware degrades ordering rather than the whole observability pipeline.🤖 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/src/api/runtime/subscriber_dispatcher.rs` around lines 812 - 822, Update the DispatcherMessage::Barrier handling to avoid blocking the dispatcher indefinitely on publications.recv(). Use a bounded receive-timeout loop that periodically logs the stalled event name and lineage, then abandons the barrier after the configured limit while still releasing permit; preserve publication handling when messages arrive and ensure all dispatcher events and flush_subscribers() can continue despite a wedged middleware future.crates/core/src/api/tool.rs (1)
828-828: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash rg -nP -C3 '\btool_(request_intercepts|conditional_execution)\s*\(' --type=rust --type=go --type=md rg -nP -C3 'tool_conditional_execution|toolConditionalExecution' docsRepository: NVIDIA/NeMo-Relay
Length of output: 155
🏁 Script executed:
#!/bin/bash set -euo pipefail # Map relevant file(s) and nearby symbols ast-grep outline crates/core/src/api/tool.rs --view expanded || true # Find definitions and uses of the async entrypoints rg -n --hidden --glob '!target' --glob '!node_modules' \ '\btool_(request_intercepts|conditional_execution)\b|toolCallWithSubscriberSnapshot|tool_conditional_execution|toolRequestIntercepts' \ crates docs python .github # Inspect the tool API section around the reported lines sed -n '780,930p' crates/core/src/api/tool.rs # Search for docs/migration coverage for async tool APIs rg -n --hidden --glob '!target' --glob '!node_modules' \ 'migration|tool_request_intercepts|tool_conditional_execution|conditional_execution|request_intercepts' \ docs crates python README.mdRepository: NVIDIA/NeMo-Relay
Length of output: 50374
🏁 Script executed:
#!/bin/bash set -euo pipefail # Narrowly inspect the tool API and nearby call sites if the first pass was sparse git ls-files crates/core/src/api/tool.rs docs | sed -n '1,120p' wc -l crates/core/src/api/tool.rs sed -n '320,430p' crates/core/src/api/tool.rs sed -n '820,900p' crates/core/src/api/tool.rsRepository: NVIDIA/NeMo-Relay
Length of output: 12433
🏁 Script executed:
#!/bin/bash set -euo pipefail python3 - <<'PY' from pathlib import Path p = Path('crates/core/src/api/tool.rs') text = p.read_text() for needle in ['pub async fn tool_request_intercepts', 'pub async fn tool_conditional_execution', 'tool_call_with_subscriber_snapshot', 'block_on_sync_ffi', 'execute_tokio_future']: print(f'## {needle}') for i, line in enumerate(text.splitlines(), 1): if needle in line: start = max(1, i-3) end = min(len(text.splitlines()), i+8) for j in range(start, end+1): print(f'{j}:{text.splitlines()[j-1]}') break else: print('not found') PYRepository: NVIDIA/NeMo-Relay
Length of output: 2256
🏁 Script executed:
#!/bin/bash set -euo pipefail # Inspect the relevant doc snippets and migration guide references sed -n '150,230p' docs/integrate-into-frameworks/code-examples.mdx sed -n '285,325p' docs/instrument-applications/code-examples.mdx sed -n '1,140p' docs/reference/migration-guides.mdx # Inspect the Rust-facing and Node/Python wrappers around these APIs sed -n '1440,1495p' crates/python/src/py_api/mod.rs sed -n '3870,3925p' crates/node/src/api/mod.rs sed -n '165,225p' crates/ffi/src/api/mod.rsRepository: NVIDIA/NeMo-Relay
Length of output: 16304
🏁 Script executed:
#!/bin/bash set -euo pipefail sed -n '150,230p' docs/integrate-into-frameworks/code-examples.mdx sed -n '285,325p' docs/instrument-applications/code-examples.mdx sed -n '1,140p' docs/reference/migration-guides.mdx sed -n '1440,1495p' crates/python/src/py_api/mod.rs sed -n '3870,3925p' crates/node/src/api/mod.rs sed -n '165,225p' crates/ffi/src/api/mod.rsRepository: NVIDIA/NeMo-Relay
Length of output: 16304
Update the Rust examples to await the async tool helpers
tool_request_interceptsandtool_conditional_executionare async now, but the Rust snippets indocs/integrate-into-frameworks/code-examples.mdxanddocs/instrument-applications/code-examples.mdxstill call them synchronously. Switch those examples to.awaitso they compile and match the current API.🤖 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/src/api/tool.rs` at line 828, Update the Rust examples in the referenced documentation to await calls to the async helpers tool_request_intercepts and tool_conditional_execution. Add .await at each call site while preserving the surrounding example logic and current API usage.Source: Path instructions
crates/core/src/plugin/dynamic/native.rs (4)
1613-1657: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
next_refleaks when the native callback panics or returns an invalid state.Both error paths reclaim only the completion
Arc(and the invocation string); theArc<NativeAsyncNext>raw pointer created at Line 1615 is never reclaimed, so the continuation plus its retained_callback_user_dataleak. Because that user data holds the plugin instance, the dynamic library can never be unloaded, and a repeatedly panicking plugin leaks without bound.🩹 Reclaim both handles on the failure paths
let state = match catch_unwind(AssertUnwindSafe(|| unsafe { cb( user_data.ptr, invocation as *const NemoRelayNativeString, next_ref .map(|next| next as *const NemoRelayNativeAsyncNext) .unwrap_or(ptr::null()), completion_ref as *const NemoRelayNativeAsyncCompletion, ) })) { Ok(state) => state, Err(_) => { unsafe { drop(Arc::from_raw( completion_ref as *const NativeAsyncCompletion, )); + if let Some(next_ref) = next_ref { + drop(Arc::from_raw(next_ref as *const NativeAsyncNext)); + } native_string_free(invocation as *mut NemoRelayNativeString); } return Err(FlowError::Internal("native async callback panicked".into())); } }; unsafe { native_string_free(invocation as *mut NemoRelayNativeString) }; let state = match NemoRelayNativeAsyncCallbackState::try_from(state) { Ok(state) => state, Err(()) => { unsafe { drop(Arc::from_raw( completion_ref as *const NativeAsyncCompletion, )); + if let Some(next_ref) = next_ref { + drop(Arc::from_raw(next_ref as *const NativeAsyncNext)); + } } return Err(FlowError::Internal( "native async callback returned an invalid state".into(), )); } };📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.let completion_ref = Arc::into_raw(completion.clone()) as usize; let next_ref = match (next, runtime) { (Some(inner), Some(runtime)) => Some(Arc::into_raw(Arc::new(NativeAsyncNext::new( inner, runtime, Some(user_data.clone()), ))) as usize), (None, None) => None, _ => unreachable!("runtime is present exactly for native async intercepts"), }; let state = match catch_unwind(AssertUnwindSafe(|| unsafe { cb( user_data.ptr, invocation as *const NemoRelayNativeString, next_ref .map(|next| next as *const NemoRelayNativeAsyncNext) .unwrap_or(ptr::null()), completion_ref as *const NemoRelayNativeAsyncCompletion, ) })) { Ok(state) => state, Err(_) => { unsafe { drop(Arc::from_raw( completion_ref as *const NativeAsyncCompletion, )); if let Some(next_ref) = next_ref { drop(Arc::from_raw(next_ref as *const NativeAsyncNext)); } native_string_free(invocation as *mut NemoRelayNativeString); } return Err(FlowError::Internal("native async callback panicked".into())); } }; unsafe { native_string_free(invocation as *mut NemoRelayNativeString) }; let state = match NemoRelayNativeAsyncCallbackState::try_from(state) { Ok(state) => state, Err(()) => { unsafe { drop(Arc::from_raw( completion_ref as *const NativeAsyncCompletion, )); if let Some(next_ref) = next_ref { drop(Arc::from_raw(next_ref as *const NativeAsyncNext)); } } return Err(FlowError::Internal( "native async callback returned an invalid state".into(), )); } };🤖 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/src/plugin/dynamic/native.rs` around lines 1613 - 1657, Reclaim the raw Arc created for next_ref in both failure branches of the native async callback flow: when catch_unwind handles a panic and when NemoRelayNativeAsyncCallbackState::try_from returns Err. Drop it with the matching NativeAsyncNext type before returning, while preserving the existing completion and invocation cleanup.
2099-2143: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
Chunk-encoding failure is reported to the plugin as a clean end of stream.
When
native_string_from_jsonreturnsNonethe loopbreaks and falls through to the terminalcb(..., ptr::null(), ptr::null(), true)call, so the plugin sees a successful completion after silently dropping the remainder of the stream. Deliver an error terminal instead.🤖 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/src/plugin/dynamic/native.rs` around lines 2099 - 2143, The native_string_from_json failure branch in the chunk-processing loop must report an error terminal instead of breaking into the successful completion callback. Update that None path to invoke cb with a non-null error message and end-of-stream flag false, perform the same cleanup via callback_guard.finish(), and return; preserve the existing successful terminal callback for fully encoded streams.
3123-3126: 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
Legacy synchronous ABI callbacks now block a runtime worker inside the returned future.
These wrappers only wrap the existing blocking FFI call in
Box::pin(async move { ... }); nothing yields. Since the middleware chains nowawaitthem on a tokio worker, a slow v1/v2 plugin stalls the reactor instead of a blocking context. Considerspawn_blocking(or the existingnative_runtime()blocking pool) for the sync-ABI path so async-ification does not change the concurrency profile of installed plugins.Also applies to: 3179-3182, 3235-3263, 3354-3357, 3369-3372, 3513-3540, 3552-3608
🤖 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/src/plugin/dynamic/native.rs` around lines 3123 - 3126, Update the synchronous legacy ABI callback wrappers around call_event_sanitize_callback and the additionally listed callback ranges to execute blocking FFI work through spawn_blocking or the established native_runtime() blocking pool before returning the future. Preserve callback arguments, result propagation, and existing async ABI behavior while ensuring v1/v2 plugin calls do not run directly on Tokio runtime workers.
3968-3988: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift
Thread-per-chunk on the stream poll path.
spawn_with_continuation_context(...).join()creates and joins an OS thread for everypoll_relay_llm_streamcall, i.e. once per streamed chunk, plus ablock_onper thread. For token-level streaming this dominates the per-chunk cost. A reusable worker (or driving the stream on the runtime and handing chunks across a channel) would remove the per-chunk spawn.🤖 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/src/plugin/dynamic/native.rs` around lines 3968 - 3988, Refactor the stream polling flow around poll_relay_llm_stream and spawn_with_continuation_context so each chunk does not create and join a new OS thread or invoke a per-call block_on. Keep the stream state and poisoned-lock error behavior intact while introducing a reusable worker, runtime-driven stream task, or channel-based mechanism that persists across poll calls and delivers chunks efficiently.crates/core/src/plugin/dynamic/worker.rs (1)
1127-1136: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Worker event sanitizers discard the incoming
fields, dropping earlier chain edits.
invoke_event_sanitizeonly sends the event, so the worker recomputes fields from the raw event and the returned value replaces whatever a higher-priority sanitizer already wrote (event_sanitize_snapshot_chainapplies the result verbatim). Forwardingfieldsin the invoke payload would make the worker surface compose like the in-process one.🤖 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/src/plugin/dynamic/worker.rs` around lines 1127 - 1136, The EventSanitizeFn closure currently ignores its incoming EventSanitizeFields, causing worker sanitizers to overwrite prior chain edits. Update the callback around invoke_event_sanitize to forward the captured fields in the invocation payload, preserving earlier sanitizer changes while keeping the existing event and callback_name handling unchanged.crates/core/src/stream.rs (1)
417-422: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Join failure discards the real terminal error.
When the finalization task panics or is cancelled, the stored
terminal_result(the collector/provider error the caller actually needs) is dropped and replaced by an internal join error. Prefer surfacing the original error and logging the join failure.♻️ Preserve `terminal_result`
Poll::Ready(Err(error)) => { this.finalization = None; - Poll::Ready(Some(Err(FlowError::Internal(format!( - "stream finalization task failed: {error}" - ))))) + log::error!( + target: "nemo_relay.runtime", + event = "stream_finalization_task_failed"; + "LLM stream END finalization task failed: {error}" + ); + match this.terminal_result.take() { + Some(result) => Poll::Ready(Some(result)), + None => Poll::Ready(Some(Err(FlowError::Internal(format!( + "stream finalization task failed: {error}" + ))))), + } }🤖 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/src/stream.rs` around lines 417 - 422, Update the Poll::Ready(Err(error)) branch of the finalization task handling to preserve and return the stored terminal_result instead of replacing it with FlowError::Internal. Log the join failure represented by error, then clear finalization while surfacing the original collector/provider error to the caller.crates/python/src/py_api/mod.rs (1)
76-92: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Where is the pyo3-async-runtimes tokio runtime configured (worker thread count)? rg -nP -C4 'tokio::(init|init_with_runtime|get_runtime)|Builder::new_multi_thread|worker_threads' --type=rust crates/pythonRepository: NVIDIA/NeMo-Relay
Length of output: 1721
🏁 Script executed:
#!/bin/bash set -euo pipefail echo '--- py_api/mod.rs around block_on_sync_middleware ---' sed -n '1,180p' crates/python/src/py_api/mod.rs echo echo '--- call sites for block_on_sync_middleware ---' rg -n "block_on_sync_middleware" crates/python/src -n -C 3 echo echo '--- explicit tokio runtime configuration in crates/python ---' rg -n "Builder::new_(multi|current)_thread|worker_threads\\(|get_runtime\\(|Handle::try_current\\(" crates/python/src -n -C 3Repository: NVIDIA/NeMo-Relay
Length of output: 10312
🌐 Web query:
pyo3_async_runtimes tokio get_runtime multi_thread worker_threads single-thread runtime documentation💡 Result:
The
pyo3-async-runtimescrate provides mechanisms to manage atokioruntime within a Python environment, primarily through thepyo3_async_runtimes::tokiomodule [1]. For configuring and retrieving the runtime, the following approaches are supported: Attribute Macro The#[pyo3_async_runtimes::tokio::main]macro allows for declarative configuration of the runtime on an async main function [2][3]. It supports the following parameters: - flavor: Selects the runtime type, either "multi_thread" or "current_thread" [2][3]. - worker_threads: Sets the number of worker threads (defaults to the number of system CPUs) [2][3]. Example: #[pyo3_async_runtimes::tokio::main(flavor = "multi_thread", worker_threads = 10)] async fn main -> PyResult<> { Ok() } Manual Initialization and Access If you need to configure the runtime manually, you can pass atokio::runtime::Builderintopyo3_async_runtimes::tokio::init()[4]. Thetokio::runtime::Runtimeis then instantiated lazily upon the first call topyo3_async_runtimes::tokio::get_runtime()[5][4]. - Custom Configuration: To use a specific configuration (e.g., setting worker threads), construct atokio::runtime::Builder, configure it, and pass it toinit[4]. - Multi-threaded: UseBuilder::new_multi_thread()and the.worker_threads()method [6][7]. - Single-threaded: UseBuilder::new_current_thread()[4]. Note that if you use the current-thread scheduler, you must manually spawn the thread that runs the runtime during initialization [4]. Example of manual initialization: let mut builder = tokio::runtime::Builder::new_current_thread; builder.enable_all; pyo3_async_runtimes::tokio::init(builder); // The runtime is retrieved lazily via get_runtime let runtime = pyo3_async_runtimes::tokio::get_runtime;Citations:
- 1: https://docs.rs/pyo3-async-runtimes/latest/pyo3_async_runtimes/tokio/index.html
- 2: https://docs.rs/pyo3-async-runtimes/latest/pyo3_async_runtimes/tokio/attr.main.html
- 3: https://docs.rs/pyo3-async-runtimes-macros/latest/pyo3_async_runtimes_macros/attr.tokio_main.html
- 4: https://docs.rs/crate/pyo3-async-runtimes/0.24.0
- 5: https://docs.rs/pyo3-async-runtimes/latest/src/pyo3_async_runtimes/tokio.rs.html
- 6: https://docs.rs/pyo3-async-runtimes/latest/pyo3_async_runtimes/tokio/re_exports/runtime/struct.Builder.html
- 7: https://docs.rs/tokio/latest/%20tokio/runtime/struct.Builder.html
🏁 Script executed:
#!/bin/bash set -euo pipefail echo '--- tests mentioning current_thread, worker_threads, or block_on_sync_middleware ---' rg -n "current_thread|worker_threads|block_on_sync_middleware|get_runtime\\(" crates/python -g '*test*' -g '*tests*' -C 2 echo echo '--- all Python-side Tokio runtime builders in the repo ---' rg -n "Builder::new_(multi|current)_thread|worker_threads\\(" crates/python crates -g '*.rs' -C 2Repository: NVIDIA/NeMo-Relay
Length of output: 28226
Add a single-worker Tokio regression test for the sync bridge (
crates/python/src/py_api/mod.rs:81-86).block_on_sync_middlewareavoids reentrancy by offloading to a scoped thread, but apyo3_async_runtimes::tokioruntime configured withworker_threads(1)can still starve the only worker while the future needs the scheduler. The current-thread coverage does not exercise that case.🤖 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/python/src/py_api/mod.rs` around lines 76 - 92, Add a regression test for block_on_sync_middleware using a Tokio runtime configured with exactly one worker thread, and execute it from within that runtime while the bridged future requires scheduler progress. Verify the call completes successfully without deadlocking or starving the worker, while preserving the existing current-thread coverage.crates/python/src/py_callable.rs (2)
62-94: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
Cancellation plumbing looks sound; one poisoning nit.
Schedule/complete/cancel ordering is correct (task cleared on
Ready,cancelledobserved by__call__afterensure_future, un-scheduled awaitables closed on drop). The.expect("scheduled awaitable")locks inpollandDropwill panic if the mutex is ever poisoned; sinceDropruns during unwinding, that turns a single failure into an abort. Preferunwrap_or_else(PoisonError::into_inner)in both places.Also applies to: 96-155
🤖 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/python/src/py_callable.rs` around lines 62 - 94, Replace the poisoning-sensitive `expect("scheduled awaitable")` mutex locks in both `CancellablePyFuture::poll` and `CancellablePyFuture::drop` with recovery using `PoisonError::into_inner`, preserving the existing task-clearing, cancellation, and completion behavior.
770-833: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Five near-identical invocation preambles; extract the shared bridge.
wrap_py_tool_fn,wrap_py_tool_conditional_fn,wrap_py_tool_request_intercept_fn,wrap_py_tool_exec_fn, and the LLM equivalents below all repeat: capture locals → copy invocation context →loop_affine_callback→ 2-or-4-armmatch (context, publication && !loop_affine)call dispatch → locals-aware split. The 4-arm publication match is duplicated verbatim in three places, which is exactly where a future fix will be applied inconsistently. A single helper taking(callback, args_tuple, publication, invocation_context, task_locals)and returning the split outcome would collapse all of them.Also applies to: 836-880, 883-910, 915-942
🤖 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/python/src/py_callable.rs` around lines 770 - 833, Extract the repeated invocation bridge from wrap_py_tool_fn, wrap_py_tool_conditional_fn, wrap_py_tool_request_intercept_fn, wrap_py_tool_exec_fn, and the equivalent LLM wrappers into one shared helper. Have it accept the callback, argument tuple, publication flag, invocation context, and task locals, perform the existing context/publication dispatch, and return the locals-aware split result. Replace each duplicated preamble and four-arm dispatch with this helper while preserving current error handling and invocation behavior.crates/python/src/py_types/observability.rs (1)
379-389: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Docs state the precondition but no longer state the consequence.
The removed sentence about re-entrant calls not establishing the delivery barrier was the only hint about what happens if a user ignores "outside subscriber and middleware callbacks".
flush_subscribersincrates/python/src/py_api/mod.rs(Lines 1565-1570) phrases the same constraint as an explicit directive; mirroring that wording here — and saying what breaks on re-entrant use — keeps the public API docs actionable.🤖 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/python/src/py_types/observability.rs` around lines 379 - 389, Update the force_flush documentation to explicitly state that it must not be called from subscriber or middleware callbacks, and describe that re-entrant calls do not establish the queued-delivery barrier. Keep the existing behavior and sink-flush description unchanged.crates/python/tests/coverage/coverage_tests.rs (1)
664-757: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Assertions match the new wrapper messages; test name is now stale.
The error-substring expectations line up with the wrapper text ("unexpected type", "callable failed", propagated
tool boom/resp boom), and thellm_req_badarity mismatch correctly exercises the call-failure branch rather than the outcome-type branch. The enclosing test is still namedtest_sync_wrapper_fallbacks_and_helperseven though the fallback-to-default behavior it covered is gone — rename to reflect error propagation.🤖 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/python/tests/coverage/coverage_tests.rs` around lines 664 - 757, Rename the enclosing test function test_sync_wrapper_fallbacks_and_helpers to reflect that it now verifies wrapper error propagation, while preserving all existing assertions and test behavior.crates/python/tests/coverage/nemo_guardrails_coverage_tests.rs (1)
318-334: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Drain logic is correct; two teardown ergonomics issues.
current_task(loop)resolves correctly becauserun_until_completewraps the coroutine in a task, and cancel-then-gather(return_exceptions=True)is the right shape. Two things:
- The
.unwrap()chain runs beforematch result { Err(payload) => resume_unwind }, so a drain failure replaces the real test panic with a teardown panic and hides the actual assertion failure.- The embedded source uses escaped
\non one line; a raw string literal with real newlines is far easier to edit.♻️ Suggested readability change
- let drain = PyModule::from_code( - py, - &CString::new( - "import asyncio\nasync def drain(loop):\n current = asyncio.current_task(loop)\n pending = asyncio.all_tasks(loop) - {current}\n for task in pending:\n task.cancel()\n if pending:\n await asyncio.gather(*pending, return_exceptions=True)\n", - ) - .unwrap(), + let drain_source = r#" +import asyncio + +async def drain(loop): + current = asyncio.current_task(loop) + pending = asyncio.all_tasks(loop) - {current} + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) +"#; + let drain = PyModule::from_code( + py, + &CString::new(drain_source).unwrap(),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.let drain_source = r#" import asyncio async def drain(loop): current = asyncio.current_task(loop) pending = asyncio.all_tasks(loop) - {current} for task in pending: task.cancel() if pending: await asyncio.gather(*pending, return_exceptions=True) "#; let drain = PyModule::from_code( py, &CString::new(drain_source).unwrap(), &CString::new("drain_test_loop.py").unwrap(), &CString::new("drain_test_loop").unwrap(), ) .unwrap() .getattr("drain") .unwrap() .call1((&event_loop,)) .unwrap(); event_loop .call_method1("run_until_complete", (drain,)) .unwrap();🤖 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/python/tests/coverage/nemo_guardrails_coverage_tests.rs` around lines 318 - 334, Update the teardown drain flow around the embedded drain module and run_until_complete call so drain failures are captured and handled after the test result, allowing any original test panic to remain the reported failure instead of being replaced by an unwrap panic. Also replace the escaped one-line Python source passed to PyModule::from_code with a multiline raw string literal while preserving the existing drain behavior.crates/python/tests/coverage/py_api_coverage_tests.rs (2)
335-346: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
conditional_allowedasserts nothing.It is set to
Trueunconditionally on the line after the awaited call, so the"conditional_allowed": trueassertion at Line 587 is vacuous. Either drop the field or derive it from the call, e.g. capture the blocked case:💚 Make the assertion meaningful
async def run_standalone(api, request): tool_args = await api.tool_request_intercepts("demo-tool", {"value": 1}) await api.tool_conditional_execution("demo-tool", tool_args) - conditional_allowed = True + try: + await api.tool_conditional_execution("demo-tool", {"value": -1}) + conditional_blocked = False + except Exception: + conditional_blocked = True llm_outcome = await api.llm_request_intercepts("demo-llm", request) await api.llm_conditional_execution(llm_outcome.request) return { "tool_value": tool_args["value"], - "conditional_allowed": conditional_allowed, + "conditional_blocked": conditional_blocked, "llm_header": llm_outcome.request.headers["x-intercepted"], }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.async def run_standalone(api, request): tool_args = await api.tool_request_intercepts("demo-tool", {"value": 1}) await api.tool_conditional_execution("demo-tool", tool_args) try: await api.tool_conditional_execution("demo-tool", {"value": -1}) conditional_blocked = False except Exception: conditional_blocked = True llm_outcome = await api.llm_request_intercepts("demo-llm", request) await api.llm_conditional_execution(llm_outcome.request) return { "tool_value": tool_args["value"], "conditional_blocked": conditional_blocked, "llm_header": llm_outcome.request.headers["x-intercepted"], }🤖 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/python/tests/coverage/py_api_coverage_tests.rs` around lines 335 - 346, Update run_standalone so conditional_allowed reflects the result of tool_conditional_execution rather than being assigned True unconditionally; capture the call’s blocked/allowed outcome and preserve the existing return field so the assertion meaningfully validates conditional execution.
928-937: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Test name promises more than it verifies.
block_on_sync_middlewaredrives the future onpyo3_async_runtimes::tokio::get_runtime(), not on thecurrent_threadruntime created here — so this only asserts that being inside some runtime routes to the scoped-thread path. It never exercises reentry into the runtime that actually owns the future. Consider asserting against the pyo3-async-runtimes runtime handle (or rename to..._offloads_when_inside_a_runtime).🤖 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/python/tests/coverage/py_api_coverage_tests.rs` around lines 928 - 937, Update synchronous_middleware_bridge_avoids_tokio_runtime_reentry to execute its async block on the pyo3_async_runtimes::tokio::get_runtime() handle, ensuring the test exercises reentry into the runtime used by block_on_sync_middleware. Preserve the existing result assertion; alternatively, rename the test to reflect only that it offloads when running inside an arbitrary Tokio runtime.crates/python/tests/coverage/py_callable_coverage_tests.rs (2)
21-90: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Duplicates the module installation already implemented in
crates/python/tests/support/mod.rs.
init_python_test_lockedincrates/python/tests/support/mod.rs(Lines 83-102) loads the sameinclude_str!("../../../../python/nemo_relay/_event_sanitizer_context.py")source and registersnemo_relayplusnemo_relay._event_sanitizer_contextintosys.modules. This helper re-implements that with save/restore semantics, and the barenemo_relaypackage it installs shadows the shared one (dropping_scope_stack_var, whichcopy_publication_invocationsilently tolerates). Move the RAII variant intotests/supportand have both call sites use it.🤖 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/python/tests/coverage/py_callable_coverage_tests.rs` around lines 21 - 90, Move InstalledContextModule and install_event_sanitizer_context_module into tests/support alongside init_python_test_locked, preserving their save/restore semantics. Update both existing call sites to use this shared RAII helper and remove the duplicate local implementation so the shared nemo_relay package setup is not shadowed.
757-933: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win
Good coverage of the new paths; cancellation is the gap.
Sync/async/raising/invalid sanitizers, the loop-affine awaitable path, the no-loop
asyncio.runfallback with a custom__await__, and continuation-context scope restoration are all exercised. Not covered: dropping aCancellablePyFuturebefore completion, which is the branch that cancels the scheduledasynciotask and closes an unscheduled awaitable (crates/python/src/py_callable.rsLines 82-94 and 146-155). A test that drops the future mid-flight and asserts the Python task ends up cancelled would protect the leak/cancel semantics.Based on coding guidelines requiring tests covering callback failure policy and scope-local cleanup for changed binding surfaces.
🤖 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/python/tests/coverage/py_callable_coverage_tests.rs` around lines 757 - 933, Add a cancellation test covering CancellablePyFuture: start a Python awaitable that signals it has begun, drop the future before completion, and assert the scheduled asyncio task finishes cancelled. Exercise the unscheduled-awaitable cleanup path as needed, using the existing Python test helpers and event-loop setup without altering current success or error tests.Source: Coding guidelines
crates/python/tests/support/mod.rs (2)
27-32: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Process-global env mutation is only serialized against tests that take the lock.
set_var/remove_varare process-wide and unsynchronized with respect to any test that does not callinit_python_test. In this PR,synchronous_middleware_bridge_avoids_tokio_runtime_reentryandexecution_next_context_restores_scoperun without the guard, so they execute concurrently with a test that is clearingNEMO_RELAY_RUNTIME_OWNER/ rewritingXDG_CONFIG_HOME. Today neither reads those, but any future runtime-init that consults env (config discovery, logging) will see torn state, and that is precisely the unsoundness theunsafemarker warns about.Cheapest fix: have every test in the crate acquire
lock_python_test(), even the Python-free ones.Lock/field ordering is correct — the
Dropbody restores env before_lockis released.Also applies to: 41-58, 60-69
🤖 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/python/tests/support/mod.rs` around lines 27 - 32, Ensure every test in the crate, including synchronous_middleware_bridge_avoids_tokio_runtime_reentry and execution_next_context_restores_scope, acquires lock_python_test() before executing. Keep the existing environment restoration and guard/drop ordering unchanged so all process-global environment mutations remain serialized.
65-69: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
Isolated config home is never created or removed.
nemo-relay-python-tests-<pid>under the temp dir is set asXDG_CONFIG_HOMEbut never created and never cleaned up, so it leaves a directory behind on any run that writes config. Not blocking; aTempDir-style cleanup inPythonTestGuard::drop(keyed to the last guard) would keep CI workspaces clean.🤖 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/python/tests/support/mod.rs` around lines 65 - 69, The Python test guard’s isolated config directory is configured but not lifecycle-managed. Update PythonTestGuard initialization and its Drop implementation to create the temp directory before setting XDG_CONFIG_HOME, then remove it when the last guard is dropped, preserving shared-guard behavior and avoiding cleanup while other guards remain active.
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
|
Moving this back to draft and holding it for 0.8. The current single-flight design avoids blocking Tokio executor threads, but concurrent subagents, parallel tool calls, or agent fan-out within one Relay runtime share a Rampart admission slot and can trigger fail-closed fallback. I want to do the shared-runtime concurrency design and validation properly before asking to land this. |
|
What we have now The managed call waits at the observability checkpoint.
Alternative concurrency implementation: rejected for this PR due to its blast radius The agent would drop off a copy of the event and continue immediately.
That alternative changes Relay's managed sanitizer and event-publication contract rather than only Rampart's scheduling. The prototype also dropped most selected bodies under burst load and accumulated detached tasks. It belongs in a separate core design if Relay wants asynchronous ordered event transformation generally. |
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com> # Conflicts: # crates/core/src/api/llm.rs # crates/core/src/api/runtime.rs # crates/core/src/api/runtime/continuation_context.rs # crates/core/src/api/runtime/state.rs # crates/core/src/api/runtime/subscriber_dispatcher.rs # crates/core/src/api/tool.rs # crates/core/src/plugin/dynamic/native.rs # crates/core/tests/integration/middleware_tests.rs # crates/core/tests/integration/pipeline_tests.rs # crates/core/tests/integration/subscriber_dispatcher_tests.rs # crates/core/tests/unit/continuation_context_tests.rs # crates/core/tests/unit/dynamic_worker_tests.rs # crates/core/tests/unit/native_plugin_tests.rs # crates/core/tests/unit/subscriber_dispatcher_tests.rs # crates/node/src/api/mod.rs # crates/node/src/callback_factory.rs # crates/node/src/promise_call.rs # crates/node/tests/llm_tests.mjs # crates/node/tests/tools_tests.mjs # crates/plugin/src/lib.rs # crates/plugin/tests/typed_callbacks.rs # crates/python/src/py_callable.rs # docs/about-nemo-relay/concepts/middleware.mdx # docs/build-plugins/dynamic-plugins/native-dynamic/about.mdx # docs/reference/event-sanitizers.mdx # docs/reference/migration-guides.mdx # python/tests/test_llm.py
Signed-off-by: Alex Fournier <afournier@nvidia.com>
|
/ok to test 9639f29 |
Signed-off-by: Alex Fournier <afournier@nvidia.com> # Conflicts: # Cargo.toml
|
/ok to test 1e37b71 |
|
Concurrency follow-up is on current head
Post-merge validation is green for pre-commit, Python ( |
4 participants
Overview
Add
pii_rampart, a separate first-party PII redaction plugin that runs the pinnednationaldesignstudio/rampartONNX model inside the Relay Rust process.The existing
pii_redactionplugin remains deterministic.pii_rampartowns its model-backed configuration and lifecycle without adding a generic inference surface or changing the gRPC worker protocol.This PR builds on the async primary-middleware boundary introduced by #571. Rampart uses that boundary to await sanitization without running tokenization or ONNX inference on Tokio executor threads or Tokio's shared blocking pool.
Details
pii_rampartcomponent kind with its own configuration, registration lifecycle, and Rust/Python/Node/Go helpers.tract-onnx. Activation requires an absolute local snapshot path and verifies SHA-256 digests for the graph, config, vocabulary, and tokenizer metadata.tract-onnxand the directrayondependency behind the crate'srampartfeature. Rayon was already present transitively through Tract; the lockfile adds no new package.Real-model concurrency validation on an Apple M4 Pro used the pinned snapshot through public managed OpenAI, Anthropic, tool, event, and streaming paths:
spawn_blockingpath reached about 501 ms p50 and failed 12/16 calls in the same condition.The main package-size cost remains unchanged. A minimal unstripped release binary grew from 4,307,520 bytes without Rampart to 36,344,480 bytes with it. Prior local package builds produced a 17,693,724-byte Python wheel and a 16,717,279-byte Node tarball.
Validation:
uv run pre-commit run --all-filescargo clippy --workspace --all-targets -- -D warningsjust test-python(616 passed)just test-node(347 passed)just test-gocargo test -p nemo-relay-pii-redaction --all-features(144 passed)cargo test -p nemo-relay-pii-redaction --no-default-features(105 passed)just test-rustpassed every compiled unit and integration suite, including 1,085 core tests and 144 PII tests. Its final core doctest failed on the pre-existingnemo_relay::Resultexample incrates/core/src/api/runtime/scope_stack.rs; the same invalid example is present onmainand is unrelated to this change.Where should the reviewer start?
Start with
crates/pii-redaction/src/rampart/mod.rsfor the independent plugin boundary, thenprefilter.rsfor the pinned model's structured-input contract,model.rsfor model ownership,tokenizer.rsfor offset fidelity, andsanitizer.rsfor the dedicated executor, bounded admission, selection, and fail-closed behavior.The key tradeoff is explicit: running in-process avoids a separate deployment and IPC path, but it adds roughly 32 MB to an unstripped release binary and gives up process-level crash isolation. Calls await sanitization for up to the bounded deadline; sustained overload favors privacy and bounded resource use by redacting or omitting selected observability content rather than growing an unbounded queue.
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)