fix: redact PII in mark and generic scope events [2/6]

crates/cli/tests/coverage/plugins_tests.rs

@@ -296,6 +296,7 @@ fn typed_editor_model_contains_pii_redaction_options() {
         &["builtin", "local_model"]
     );
     assert_eq!(schema.field("codec").unwrap().kind, EditorFieldKind::Enum);
+    assert_eq!(schema.field("mark").unwrap().kind, EditorFieldKind::Boolean);
     assert_eq!(
         schema.field("tool_output").unwrap().kind,
         EditorFieldKind::Boolean

crates/pii-redaction/Cargo.toml

@@ -26,6 +26,8 @@ sha2 = "0.11"
 schemars = { version = "0.8", optional = true }
 
 [dev-dependencies]
-nemo-relay = { workspace = true, features = ["openinference"] }
+nemo-relay = { workspace = true, features = ["openinference", "otel"] }
 futures = "0.3"
 tokio = { version = "1", features = ["rt", "macros", "sync", "test-util", "rt-multi-thread", "time"] }
+opentelemetry_sdk = { workspace = true, features = ["trace", "testing"] }
+tempfile = "3"

crates/pii-redaction/src/builtin.rs

@@ -9,8 +9,11 @@ use serde::de::DeserializeOwned;
 use serde_json::Value as Json;
 use sha2::{Digest, Sha256};
 
+use nemo_relay::api::event::{CategoryProfile, Event};
 use nemo_relay::api::llm::LlmRequest;
-use nemo_relay::api::runtime::{LlmSanitizeRequestFn, LlmSanitizeResponseFn, ToolSanitizeFn};
+use nemo_relay::api::runtime::{
+    EventSanitizeFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, ToolSanitizeFn,
+};
 use nemo_relay::codec::anthropic::AnthropicMessagesCodec;
 use nemo_relay::codec::openai_chat::OpenAIChatCodec;
 use nemo_relay::codec::openai_responses::OpenAIResponsesCodec;
@@ -256,6 +259,29 @@ pub(super) fn tool_sanitize_callback(backend: CompiledBuiltinBackend) -> ToolSan
     Arc::new(move |_name: &str, payload: Json| backend.sanitize_json_preorder_dfs(payload))
 }
 
+pub(super) fn event_sanitize_callback(backend: CompiledBuiltinBackend) -> EventSanitizeFn {
+    Arc::new(move |event, mut fields| {
+        if matches!(event, Event::Scope(_))
+            && event
+                .category()
+                .is_some_and(|category| matches!(category.as_str(), "tool" | "llm"))
+        {
+            return fields;
+        }
+
+        fields.data = fields
+            .data
+            .map(|data| backend.sanitize_json_preorder_dfs(data));
+        fields.metadata = fields
+            .metadata
+            .map(|metadata| backend.sanitize_json_preorder_dfs(metadata));
+        fields.category_profile = fields.category_profile.and_then(|profile| {
+            sanitize_serializable_with_backend::<CategoryProfile>(&backend, profile).ok()
+        });
+        fields
+    })
+}
+
 pub(super) fn llm_sanitize_request_callback(
     backend: CompiledBuiltinBackend,
 ) -> LlmSanitizeRequestFn {

crates/pii-redaction/src/component.rs

@@ -81,6 +81,9 @@ pub struct PiiRedactionConfig {
     /// Whether to sanitize managed LLM response payloads.
     #[serde(default = "default_true")]
     pub output: bool,
+    /// Whether to sanitize mark event observability fields.
+    #[serde(default = "default_true")]
+    pub mark: bool,
     /// Whether to sanitize managed tool request payloads.
     #[serde(default = "default_true")]
     pub tool_input: bool,
@@ -112,6 +115,7 @@ impl Default for PiiRedactionConfig {
             mode: default_mode(),
             input: true,
             output: true,
+            mark: true,
             tool_input: true,
             tool_output: true,
             priority: default_priority(),
@@ -199,6 +203,7 @@ nemo_relay::editor_config! {
         },
         input => { label: "input", kind: Boolean },
         output => { label: "output", kind: Boolean },
+        mark => { label: "mark", kind: Boolean },
         tool_input => { label: "tool_input", kind: Boolean },
         tool_output => { label: "tool_output", kind: Boolean },
         priority => { label: "priority", kind: Integer },
@@ -423,6 +428,7 @@ fn validate_pii_redaction_plugin_config(
             "mode",
             "input",
             "output",
+            "mark",
             "tool_input",
             "tool_output",
             "priority",
@@ -497,7 +503,7 @@ fn validate_surface_selection(
     policy: &ConfigPolicy,
     config: &PiiRedactionConfig,
 ) {
-    if config.input || config.output || config.tool_input || config.tool_output {
+    if config.input || config.output || config.mark || config.tool_input || config.tool_output {
         return;
     }
 
@@ -725,6 +731,14 @@ fn register_builtin_backend(
     })?;
     let compiled = CompiledBuiltinBackend::new(builtin, config.codec.clone())?;
 
+    if config.mark {
+        ctx.register_mark_sanitize_guardrail(
+            "mark",
+            config.priority,
+            super::builtin::event_sanitize_callback(compiled.clone()),
+        )?;
+    }
+
     if config.tool_input {
         let sanitizer = tool_sanitize_callback(compiled.clone());
         ctx.register_tool_sanitize_request_guardrail("tool_input", config.priority, sanitizer)?;
@@ -734,12 +748,28 @@ fn register_builtin_backend(
         ctx.register_tool_sanitize_response_guardrail("tool_output", config.priority, sanitizer)?;
     }
     if config.input {
-        let sanitizer = llm_sanitize_request_callback(compiled.clone());
-        ctx.register_llm_sanitize_request_guardrail("input", config.priority, sanitizer)?;
+        ctx.register_llm_sanitize_request_guardrail(
+            "input",
+            config.priority,
+            llm_sanitize_request_callback(compiled.clone()),
+        )?;
+        ctx.register_scope_sanitize_start_guardrail(
+            "input",
+            config.priority,
+            super::builtin::event_sanitize_callback(compiled.clone()),
+        )?;
     }
     if config.output {
-        let sanitizer = llm_sanitize_response_callback(compiled);
-        ctx.register_llm_sanitize_response_guardrail("output", config.priority, sanitizer)?;
+        ctx.register_llm_sanitize_response_guardrail(
+            "output",
+            config.priority,
+            llm_sanitize_response_callback(compiled.clone()),
+        )?;
+        ctx.register_scope_sanitize_end_guardrail(
+            "output",
+            config.priority,
+            super::builtin::event_sanitize_callback(compiled),
+        )?;
     }
 
     Ok(())