-
Notifications
You must be signed in to change notification settings - Fork 67
fix: harden trajectory context redaction #991
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -16,15 +16,18 @@ use nemo_relay::api::runtime::{ | |||||||||||||
| BuiltinLlmCodec, EventSanitizeFn, LlmCodecIdentity, LlmSanitizeRequestFn, | ||||||||||||||
| LlmSanitizeResponseFn, ToolSanitizeFn, | ||||||||||||||
| }; | ||||||||||||||
| use nemo_relay::codec::request::AnnotatedLlmRequest; | ||||||||||||||
| use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; | ||||||||||||||
| use nemo_relay::codec::resolve::{ | ||||||||||||||
| ProviderSurface, detect_response_surface, request_codec as build_request_codec, | ||||||||||||||
| response_codec as build_response_codec, | ||||||||||||||
| }; | ||||||||||||||
| use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; | ||||||||||||||
| use nemo_relay::plugin::{PluginError, Result as PluginResult}; | ||||||||||||||
|
|
||||||||||||||
| use super::component::{BuiltinBackendConfig, validate_metric_string_attribute_allowlist}; | ||||||||||||||
| use super::component::{ | ||||||||||||||
| BuiltinBackendConfig, DEFAULT_CUSTOM_MARK_PAYLOAD_POLICY, | ||||||||||||||
| validate_metric_string_attribute_allowlist, | ||||||||||||||
| }; | ||||||||||||||
| use super::detectors::BuiltinDetector; | ||||||||||||||
| use super::overlay::BuiltinCodecName; | ||||||||||||||
| use super::trajectory::{CustomMarkPayloadPolicy, TrajectorySanitizer, is_relay_metric_mark}; | ||||||||||||||
|
|
@@ -300,7 +303,9 @@ impl CompiledBuiltinBackend { | |||||||||||||
| } | ||||||||||||||
| None => None, | ||||||||||||||
| }; | ||||||||||||||
| if trajectory.is_none() && config.custom_mark_payload_policy != "preserve" { | ||||||||||||||
| if trajectory.is_none() | ||||||||||||||
| && config.custom_mark_payload_policy != DEFAULT_CUSTOM_MARK_PAYLOAD_POLICY | ||||||||||||||
| { | ||||||||||||||
| return Err(PluginError::InvalidConfig( | ||||||||||||||
| "builtin.custom_mark_payload_policy requires builtin.preset = 'trajectory_context'" | ||||||||||||||
| .to_string(), | ||||||||||||||
|
|
@@ -638,6 +643,73 @@ impl CompiledBuiltinBackend { | |||||||||||||
| }) | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| fn sanitize_trajectory_request_with_codec( | ||||||||||||||
| &self, | ||||||||||||||
| trajectory: &TrajectorySanitizer, | ||||||||||||||
| codec: &dyn LlmCodec, | ||||||||||||||
| request: &LlmRequest, | ||||||||||||||
| ) -> Option<LlmRequest> { | ||||||||||||||
| let annotated = codec.decode(request).ok()?; | ||||||||||||||
| let sanitized = trajectory.sanitize_annotated_request(annotated)?; | ||||||||||||||
| if let Ok(request) = codec.encode(&sanitized, request) { | ||||||||||||||
| return Some(request); | ||||||||||||||
| } | ||||||||||||||
| self.encode_trajectory_request_items_incrementally(codec, request, &sanitized) | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| fn encode_trajectory_request_items_incrementally( | ||||||||||||||
| &self, | ||||||||||||||
| codec: &dyn LlmCodec, | ||||||||||||||
| request: &LlmRequest, | ||||||||||||||
| sanitized: &AnnotatedLlmRequest, | ||||||||||||||
| ) -> Option<LlmRequest> { | ||||||||||||||
| let mut output = request.clone(); | ||||||||||||||
| for (index, message) in sanitized.messages.iter().enumerate() { | ||||||||||||||
| if let Some(MessageContent::Parts(parts)) = message_content(message) { | ||||||||||||||
| for (part_index, part) in parts.iter().enumerate() { | ||||||||||||||
| let mut current = codec.decode(&output).ok()?; | ||||||||||||||
| let current_message = current.messages.get_mut(index)?; | ||||||||||||||
| let MessageContent::Parts(current_parts) = | ||||||||||||||
| message_content_mut(current_message)? | ||||||||||||||
| else { | ||||||||||||||
| return None; | ||||||||||||||
| }; | ||||||||||||||
| current_parts.get_mut(part_index)?.clone_from(part); | ||||||||||||||
| output = codec.encode(¤t, &output).ok()?; | ||||||||||||||
| } | ||||||||||||||
| } | ||||||||||||||
| if let Message::Assistant { | ||||||||||||||
| tool_calls: Some(tool_calls), | ||||||||||||||
| .. | ||||||||||||||
| } = message | ||||||||||||||
| { | ||||||||||||||
| for (call_index, call) in tool_calls.iter().enumerate() { | ||||||||||||||
| let mut current = codec.decode(&output).ok()?; | ||||||||||||||
| let Message::Assistant { | ||||||||||||||
| tool_calls: Some(current_calls), | ||||||||||||||
| .. | ||||||||||||||
| } = current.messages.get_mut(index)? | ||||||||||||||
| else { | ||||||||||||||
| return None; | ||||||||||||||
| }; | ||||||||||||||
| current_calls.get_mut(call_index)?.clone_from(call); | ||||||||||||||
| output = codec.encode(¤t, &output).ok()?; | ||||||||||||||
| } | ||||||||||||||
| } | ||||||||||||||
| let mut current = codec.decode(&output).ok()?; | ||||||||||||||
| current.messages[index] = message.clone(); | ||||||||||||||
| output = codec.encode(¤t, &output).ok()?; | ||||||||||||||
|
Comment on lines
+699
to
+701
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win Guard the message index before assignment. Line 700 indexes This function runs only after the whole-request 🐛 Proposed fix to fail closed instead of panicking let mut current = codec.decode(&output).ok()?;
- current.messages[index] = message.clone();
+ current.messages.get_mut(index)?.clone_from(message);
output = codec.encode(¤t, &output).ok()?;📝 Committable suggestion
Suggested change
🤖 Prompt for AI Agents |
||||||||||||||
| } | ||||||||||||||
|
Comment on lines
+667
to
+702
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift Bound the incremental re-encode work. Each inner iteration performs a full The cost scales with attacker- or client-controlled request size, and the fallback triggers precisely on the payloads that already failed a single whole-request encode. Consider decoding once, applying every sanitized item to that single 🤖 Prompt for AI Agents |
||||||||||||||
| if let Some(tools) = sanitized.tools.as_ref() { | ||||||||||||||
| for (index, tool) in tools.iter().enumerate() { | ||||||||||||||
| let mut current = codec.decode(&output).ok()?; | ||||||||||||||
| current.tools.as_mut()?.get_mut(index)?.clone_from(tool); | ||||||||||||||
| output = codec.encode(¤t, &output).ok()?; | ||||||||||||||
| } | ||||||||||||||
| } | ||||||||||||||
| codec.encode(sanitized, &output).ok() | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| fn sanitize_request_target_paths_incrementally( | ||||||||||||||
| &self, | ||||||||||||||
| codec: &dyn LlmCodec, | ||||||||||||||
|
|
@@ -758,6 +830,21 @@ impl CompiledBuiltinBackend { | |||||||||||||
| .then_some(payload) | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| fn sanitize_trajectory_response_with_codec( | ||||||||||||||
| &self, | ||||||||||||||
| trajectory: &TrajectorySanitizer, | ||||||||||||||
| codec: &dyn LlmResponseCodec, | ||||||||||||||
| surface: ProviderSurface, | ||||||||||||||
| payload: Json, | ||||||||||||||
| ) -> Option<Json> { | ||||||||||||||
| let annotated = codec.decode_response(&payload).ok()?; | ||||||||||||||
| let sanitized = trajectory.sanitize_annotated_response(annotated)?; | ||||||||||||||
| Some( | ||||||||||||||
| BuiltinCodecName::from_provider_surface(surface) | ||||||||||||||
| .overlay_response_payload(payload, &sanitized), | ||||||||||||||
|
Comment on lines
+840
to
+844
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🔒 Security & Privacy | 🟠 Major | ⚡ Quick win 🧩 Analysis chain🏁 Script executed: #!/bin/bash
# Description: Determine whether overlay_response_payload rewrites all choices/candidates or only the first.
set -euo pipefail
echo "===== overlay_response_payload definition ====="
ast-grep run --pattern 'fn overlay_response_payload($$$) { $$$ }' --lang rust crates 2>/dev/null \
|| fd -e rs . crates -x rg -n -A 80 'fn overlay_response_payload' {}
echo "===== indexing of choices/candidates in the overlay and codecs ====="
fd -e rs . crates/core/src/codec -x rg -n 'choices|candidates' {} | rg -n '\[0\]|first\(|get\(0\)|iter_mut|enumerate' || true
echo "===== existing multi-choice coverage for the trajectory preset ====="
fd -e rs . crates/pii-redaction -x rg -n -B 4 -A 12 'trajectory' {} | rg -n 'choices|candidates' || echo "no trajectory test exercises multi-choice payloads"Repository: NVIDIA/NeMo-Relay Length of output: 3423 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "===== all response overlay implementations ====="
sed -n '1,240p' crates/core/src/codec.rs 2>/dev/null || true
fd -e rs . crates -x rg -n -A 100 -B 5 'fn overlay_(openai_chat|openai_responses|anthropic|oci_genai|gemini)_response' {} 2>/dev/null
echo "===== response sanitization guards and trajectory path ====="
fd -e rs . crates/pii-redaction -x rg -n -A 45 -B 15 \
'sanitize_response_with_codec|sanitize_trajectory_response_with_codec|targets_normalized_single_projected_response|sanitize_annotated_response' {} 2>/dev/null
echo "===== codec response models and multi-item handling ====="
fd -e rs . crates -x rg -n -A 35 -B 10 \
'struct AnnotatedLlmResponse|struct LlmResponse|choices:|candidates:|overlay_.*parts|first_mut\(\)' {} 2>/dev/nullRepository: NVIDIA/NeMo-Relay Length of output: 10871 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "===== all response overlay implementations ====="
fd -e rs . crates -x rg -n -A 100 -B 5 \
'fn overlay_(openai_chat|openai_responses|anthropic|oci_genai|gemini)_response' {} 2>/dev/null
echo "===== response sanitization guards and trajectory path ====="
fd -e rs . crates/pii-redaction -x rg -n -A 45 -B 15 \
'sanitize_response_with_codec|sanitize_trajectory_response_with_codec|targets_normalized_single_projected_response|sanitize_annotated_response' {} 2>/dev/null
echo "===== codec response models and multi-item handling ====="
fd -e rs . crates -x rg -n -A 35 -B 10 \
'struct AnnotatedLlmResponse|struct LlmResponse|choices:|candidates:|overlay_.*parts|first_mut\(\)' {} 2>/dev/nullRepository: NVIDIA/NeMo-Relay Length of output: 10871 🏁 Script executed: #!/bin/bash
set -euo pipefail
echo "===== trajectory response path ====="
sed -n '820,930p' crates/pii-redaction/src/builtin.rs
sed -n '1035,1095p' crates/pii-redaction/src/builtin.rs
echo "===== legacy response guard ====="
rg -n -A 70 -B 15 \
'fn sanitize_response_with_codec|targets_normalized_single_projected_response' \
crates/pii-redaction/src crates/core 2>/dev/null || true
echo "===== response decoder projection ====="
rg -n -A 45 -B 15 \
'fn decode_response|AnnotatedLlmResponse|choices.*first|candidates.*first' \
crates 2>/dev/null | head -n 300Repository: NVIDIA/NeMo-Relay Length of output: 43991 Sensitive Data Exposure (CWE-201) Reachability: External · Exploitability: Moderate Fail closed on multi-choice and multi-candidate trajectory responses. The trajectory branch overlays only the first OpenAI Chat choice or Gemini candidate. Later entries remain in the original payload and can contain unsanitized model output. Add a multi-item guard before decoding, or sanitize every choice and candidate, with regression tests. 🤖 Prompt for AI Agents |
||||||||||||||
| ) | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| fn normalized_response_targets_match( | ||||||||||||||
| target_paths: &[Vec<String>], | ||||||||||||||
| annotated: &Json, | ||||||||||||||
|
|
@@ -780,6 +867,34 @@ impl CompiledBuiltinBackend { | |||||||||||||
| } | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| fn message_content(message: &Message) -> Option<&MessageContent> { | ||||||||||||||
| match message { | ||||||||||||||
| Message::System { content, .. } | ||||||||||||||
| | Message::User { content, .. } | ||||||||||||||
| | Message::Developer { content, .. } | ||||||||||||||
| | Message::Tool { content, .. } => Some(content), | ||||||||||||||
| Message::Assistant { content, .. } => content.as_ref(), | ||||||||||||||
| Message::Function { .. } | ||||||||||||||
| | Message::ToolCallItem { .. } | ||||||||||||||
| | Message::ToolResultItem { .. } | ||||||||||||||
| | Message::ProviderNative { .. } => None, | ||||||||||||||
| } | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| fn message_content_mut(message: &mut Message) -> Option<&mut MessageContent> { | ||||||||||||||
| match message { | ||||||||||||||
| Message::System { content, .. } | ||||||||||||||
| | Message::User { content, .. } | ||||||||||||||
| | Message::Developer { content, .. } | ||||||||||||||
| | Message::Tool { content, .. } => Some(content), | ||||||||||||||
| Message::Assistant { content, .. } => content.as_mut(), | ||||||||||||||
| Message::Function { .. } | ||||||||||||||
| | Message::ToolCallItem { .. } | ||||||||||||||
| | Message::ToolResultItem { .. } | ||||||||||||||
| | Message::ProviderNative { .. } => None, | ||||||||||||||
| } | ||||||||||||||
| } | ||||||||||||||
|
|
||||||||||||||
| pub(super) fn tool_sanitize_callback(backend: CompiledBuiltinBackend) -> ToolSanitizeFn { | ||||||||||||||
| let backend = Arc::new(backend); | ||||||||||||||
| Arc::new(move |_name: String, payload: Json| { | ||||||||||||||
|
|
@@ -889,8 +1004,28 @@ pub(super) fn llm_sanitize_request_callback( | |||||||||||||
| .as_object() | ||||||||||||||
| .cloned() | ||||||||||||||
| .unwrap_or_default(); | ||||||||||||||
| request.content = trajectory.sanitize_provider_payload(request.content); | ||||||||||||||
| return Ok(Some(request)); | ||||||||||||||
| let resolved = context.resolve_codec(); | ||||||||||||||
| let fallback = if resolved.is_none() { | ||||||||||||||
| backend | ||||||||||||||
| .selected_surface(context.codec()) | ||||||||||||||
| .map(build_request_codec) | ||||||||||||||
| } else { | ||||||||||||||
| None | ||||||||||||||
| }; | ||||||||||||||
| let Some(codec) = resolved.as_deref().or(fallback.as_deref()) else { | ||||||||||||||
| request.content = trajectory.sanitize_provider_payload(request.content); | ||||||||||||||
| return Ok(Some(request)); | ||||||||||||||
| }; | ||||||||||||||
| let sanitized = | ||||||||||||||
| backend.sanitize_trajectory_request_with_codec(trajectory, codec, &request); | ||||||||||||||
| if sanitized.is_none() { | ||||||||||||||
| log_llm_payload_omitted( | ||||||||||||||
| "request", | ||||||||||||||
| context.codec(), | ||||||||||||||
| "codec decode, typed sanitize, or encode failure", | ||||||||||||||
| ); | ||||||||||||||
| } | ||||||||||||||
| return Ok(sanitized); | ||||||||||||||
| } | ||||||||||||||
| request.headers = backend.sanitize_request_headers(request.headers); | ||||||||||||||
| if backend.target_path_matcher.is_empty() { | ||||||||||||||
|
|
@@ -930,7 +1065,28 @@ pub(super) fn llm_sanitize_response_callback( | |||||||||||||
| let backend = Arc::clone(&backend); | ||||||||||||||
| Box::pin(async move { | ||||||||||||||
| if let Some(trajectory) = backend.trajectory.as_ref() { | ||||||||||||||
| return Ok(Some(trajectory.sanitize_provider_payload(payload))); | ||||||||||||||
| let Some(surface) = backend.selected_surface(context.codec()) else { | ||||||||||||||
| return Ok(Some(trajectory.sanitize_provider_payload(payload))); | ||||||||||||||
| }; | ||||||||||||||
| let resolved = context.resolve_codec(); | ||||||||||||||
| let fallback = if resolved.is_none() { | ||||||||||||||
| Some(build_response_codec(surface)) | ||||||||||||||
| } else { | ||||||||||||||
| None | ||||||||||||||
| }; | ||||||||||||||
| let Some(codec) = resolved.as_deref().or(fallback.as_deref()) else { | ||||||||||||||
| return Ok(Some(trajectory.sanitize_provider_payload(payload))); | ||||||||||||||
| }; | ||||||||||||||
| let sanitized = backend | ||||||||||||||
| .sanitize_trajectory_response_with_codec(trajectory, codec, surface, payload); | ||||||||||||||
| if sanitized.is_none() { | ||||||||||||||
| log_llm_payload_omitted( | ||||||||||||||
| "response", | ||||||||||||||
| context.codec(), | ||||||||||||||
| "codec decode, typed sanitize, or encode failure", | ||||||||||||||
| ); | ||||||||||||||
| } | ||||||||||||||
| return Ok(sanitized); | ||||||||||||||
| } | ||||||||||||||
| if backend.target_path_matcher.is_empty() { | ||||||||||||||
| return Ok(Some(backend.sanitize_json_preorder_dfs(payload))); | ||||||||||||||
|
|
||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Both documents claim that unapproved numbers become
0, and no code path does that. The trajectory sanitizer never zeroes a number:sanitize_api_specific_requestandsanitize_api_specific_responsepass numeric fields through unchanged,sanitize_usagepreserves every token count, andsanitize_costpreservestotal,input, andoutput. The PR's own tests assertfrequency_penalty == 0.4,n == 2,seed == 7,cost.input == Some(0.3), andcost.output == Some(0.12). The boolean half of the sentence is correct. Drop the numeric claim at both sites and state that typed numeric fields are retained as analytics.crates/pii-redaction/README.md#L147-L147: remove "unapproved numbers become0" and state that typed numeric fields, such as token counts, cost amounts, and request tuning values, are retained as analytics.docs/configure-plugins/pii-redaction/configuration.mdx#L263-L263: apply the same correction to the identical sentence in the Trajectory Context Preset section.📍 Affects 2 files
crates/pii-redaction/README.md#L147-L147(this comment)docs/configure-plugins/pii-redaction/configuration.mdx#L263-L263🤖 Prompt for AI Agents