feat: add in-process Rampart PII redaction plugin - #558
Open
afourniernv wants to merge 86 commits into
Open
Conversation
|
Note Currently processing new changes in this PR. This may take a few minutes, please wait... ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Enterprise Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (22)
WalkthroughChangesRampart PII support now includes an ONNX detector, structured prefilter, bounded asynchronous sanitizer, plugin registration, CLI editor integration, and Node, Go, and Python APIs. Host initialization paths register the component and expose validation and configuration helpers. Rampart PII plugin
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant Host
participant PluginRegistry
participant RampartDetector
participant RampartSanitizer
participant Guardrail
Host->>PluginRegistry: register Rampart PII component
PluginRegistry->>RampartDetector: verify and load model artifacts
PluginRegistry->>RampartSanitizer: create sanitizer
PluginRegistry->>Guardrail: register surface callbacks
Guardrail->>RampartSanitizer: sanitize selected payload
RampartSanitizer->>RampartDetector: detect text batch
RampartDetector-->>RampartSanitizer: return scored spans
RampartSanitizer-->>Guardrail: return redacted payload
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
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>
afourniernv
force-pushed
the
feat/pii-worker-provider
branch
from
July 27, 2026 00:03
b73ac66 to
385c241
Compare
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
afourniernv
marked this pull request as ready for review
July 27, 2026 18:39
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>
Contributor
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.
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
afourniernv
marked this pull request as draft
July 30, 2026 02:17
Contributor
Author
|
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. |
2 tasks
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
afourniernv
changed the base branch from
wkk_relay-509-async-primary-middleware
to
main
July 31, 2026 03:07
This was referenced Jul 31, 2026
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Contributor
Author
|
/ok to test 9639f29 |
Signed-off-by: Alex Fournier <afournier@nvidia.com> # Conflicts: # Cargo.toml
Contributor
Author
|
/ok to test 1e37b71 |
2 tasks
Contributor
Author
Current-head production benchmarkI reran the Rampart validation on commit Reliability and executor behavior
Agent-shaped workloadsThese timings include request and response sanitization plus simulated provider and tool delay.
A 30-call provider lifecycle run completed without event loss at 24.25 ms p99. An 80-call streaming run, 32 mark events, and 64 manual tool lifecycle events also completed without failures or canary leaks. Selected-field size scalingThis uses a complete managed tool lifecycle, so each call sanitizes both input and output.
Admission stayed lossless through 16 concurrent 1 KiB calls and through six concurrent 8 KiB calls. At deliberate overload, 48 of 144 1 KiB bodies failed closed at concurrency 24, and 36 of 72 8 KiB bodies failed closed at concurrency 12. That is the intended bounded-admission behavior: application calls still complete, while affected observable bodies are omitted rather than queued without limit. Sustained saturationA 1,200-call soak completed in 62.04 seconds at 19.34 calls/s:
The normal Claude-shaped run peaked at 86.6 MB RSS; deliberate saturation peaked around 245-248 MB. Redaction qualityI evaluated 210 balanced English, Spanish, French, German, Italian, Portuguese, and Dutch rows from the AI4Privacy PII Masking Nano 1K dataset, using the pinned Rampart model:
This is not a replacement for the model author's evaluation, but it exercises the exact Relay preprocessing, offset projection, policy, and replacement path. Conclusion and remaining limitsThe current design is production-defensible as an opt-in sanitizer for narrow, explicitly selected observability fields, especially at or below roughly 1 KiB. The dedicated bounded Rayon pool prevents Tokio executor starvation and avoids coupling inference to Tokio's blocking pool. Admission, CPU use, and memory growth remain bounded, and overload removes observable content without changing the application payload. The main limits are also visible in the numbers: 8-16 KiB selected fields add second-scale latency, sustained CPU saturation can intentionally fail closed, and pressure runs reach roughly 248 MB RSS. I would document/select conservative payload limits rather than position this as a generic large-field sanitizer. Linux x86_64 and aarch64 performance validation also remains outstanding before making broader platform claims. |
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
afourniernv
marked this pull request as ready for review
August 3, 2026 17:58
Contributor
Author
|
/ok to test 121f563 |
Member
|
/ok to test 207fa57 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
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)
Summary by CodeRabbit
New Features
Documentation
Tests