feat(plugin): add opt-in Rampart PII redaction - #558
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughChangesRampart PII support adds 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 configuration validation helpers. Rampart PII plugin
Estimated code review effort: 5 (Critical) | ~120 minutes 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 | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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>
b73ac66 to
385c241
Compare
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
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.
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Current-head production benchmarkReran the Rampart validation on commit This head adds a bounded, content-addressed decision cache and field-level model-budget fallback:
Real Claude Code and Codex trafficI rebuilt the actual Relay CLI and ran cold direct calls plus tool-heavy sessions through Anthropic Messages and OpenAI Responses at the default four-window budget.
For comparison, the equivalent pre-cache harness retained only 1/5 Anthropic request bodies and 0/11 OpenAI Responses request bodies even with a 16-window budget. Synthetic production matrixThe 16-window run gives the direct comparison with the prior benchmark.
The soak retained all 2,400 LLM bodies with zero canary leaks and zero fail-closed bodies. Tokio heartbeat p99 was 0.449 ms, heartbeat max was 0.901 ms, and max RSS was 231.3 MB. Repeated-field timings separate the first cold inference from later exact cache hits:
At the default four-window budget, the 8 KiB field took 434.048 ms cold. The 16 KiB field was rejected field-by-field in 1.506 ms rather than spending second-scale inference time or removing the full envelope. Redaction quality stayed unchanged on 210 balanced AI4Privacy rows across seven languages:
Validation on this head:
The practical limit remains explicit: a unique cold field that exceeds its configured model-window budget is replaced in observability. Repeated coding-agent instructions, tool schemas, and history no longer repeatedly consume that budget or cause the whole request body to disappear. |
|
/ok to test 1866014 |
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Signed-off-by: Alex Fournier <afournier@nvidia.com>
|
/ok to test be1725e |
Signed-off-by: Alex Fournier <afournier@nvidia.com>
|
/ok to test 218b7eb |
License DiffCompared against Lockfile license changesLockfile License ChangesRustAdded
Removed
Updated/Changed
|
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, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
Signed-off-by: Alex Fournier <afournier@nvidia.com>
Overview
Add
pii_rampartas an opt-in, in-process native Rust plugin. The plugin runs the pinnednationaldesignstudio/rampartONNX model through Tract and sanitizes copied observability payloads without modifying provider or tool inputs and outputs.Rampart is not linked into the Relay CLI, language bindings, or default host artifacts. The existing deterministic
pii_redactioncomponent remains unchanged and autoloaded as before. Users who do not install and configure this native plugin do not carry Tract or the model runtime.Details
cdylibunderplugins/pii-rampartwith its own manifest, JSON Schema, lockfile, README, tests, and plugin lifecycle.tract-onnx, Rayon, and all Rampart code behind thenemo-relay-pii-redaction/rampartfeature used only by the standalone plugin.Validation on current
main:uv run pre-commit run --all-filescargo clippy --workspace --all-targets -- -D warningscargo clippy --locked --manifest-path plugins/pii-rampart/Cargo.toml --all-targets -- -D warningsjust test-rust(full clean rerun passed)cargo test -p nemo-relay-pii-redaction --features rampart(180 passed)cargo test --locked --manifest-path plugins/pii-rampart/Cargo.tomlThe real-model dynamic-load test exercises the built
cdylibthroughPluginHostActivation, managed tool, OpenAI, and Gemini calls. It verifies that application payloads remain unchanged, observable payloads are redacted without safe-body omission, eight concurrent calls complete, cancellation resolves, and plugin teardown completes within two seconds.The default Relay CLI dependency graph does not contain
tract-onnx. The standalone unstripped macOS plugin is approximately 28.5 MB (approximately 23.0 MB afterstrip -x); that cost is paid only when the plugin artifact is installed.Where should the reviewer start?
Start with
plugins/pii-rampart/src/lib.rsfor the native plugin and async completion lifecycle, thencrates/pii-redaction/src/rampart/sanitizer.rsfor bounded scheduling and fail-closed behavior.plugins/pii-rampart/tests/native_plugin.rsis the end-to-end dynamic-load test.The main design decision is the distribution boundary: model inference stays inside the Relay Rust process, as requested, while the optional native plugin keeps Tract out of every default Relay artifact. The follow-up PR #737 handles platform-specific archives and release publication without widening this runtime PR.
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)