Skip to main content

agent_sdk/observability/
payload.rs

1//! `GenAI` payload conversion from SDK types to semconv JSON.
2//!
3//! Converts SDK [`ChatRequest`] / [`ChatResponse`] types into the JSON
4//! schemas defined by the `GenAI` semantic conventions for
5//! `gen_ai.system_instructions`, `gen_ai.input.messages`, and
6//! `gen_ai.output.messages`.
7//!
8//! # Privacy
9//!
10//! This layer sits directly in front of observability egress —
11//! whatever shape we emit here flows into `OTel` spans, `ObservabilityStore`
12//! implementations, and ultimately into third-party collectors
13//! (Langfuse, Datadog, honeycomb, …). [`PayloadRedactor`] wraps a
14//! pluggable [`PiiDetector`] and masks PII in every text and JSON
15//! string leaf before they leave this module, so the two conversion
16//! paths (directly serialized onto spans, or handed to the store)
17//! share one audited redaction pass.
18//!
19//! Callers that do not need redaction — including existing call sites
20//! that have not adopted the detector yet — should use
21//! [`PayloadRedactor::noop`]. Financial and other PII-sensitive
22//! integrations should construct a redactor with
23//! [`agent_sdk_foundation::privacy::BaselineDetector`] or a custom
24//! detector. Routing every conversion through the [`PayloadRedactor`]
25//! type keeps the unredacted-egress surface to a single audited path.
26
27use crate::llm::{ChatRequest, ChatResponse, Content, ContentBlock, Message, Role};
28use agent_sdk_foundation::privacy::{NoopDetector, PiiDetector, mask_spans};
29use serde_json::{Value, json};
30use std::sync::Arc;
31
32use super::attrs::finish_reason_str;
33
34/// Redacts PII from payloads before they flow to observability sinks.
35///
36/// Wraps a [`PiiDetector`] and applies it to every text and JSON
37/// string leaf emitted by [`convert_system_instructions`](Self::convert_system_instructions),
38/// [`convert_input_messages`](Self::convert_input_messages), and
39/// [`convert_output_messages`](Self::convert_output_messages).
40///
41/// Construct with [`PayloadRedactor::new`] wrapping any
42/// `Arc<dyn PiiDetector>`, or with [`PayloadRedactor::noop`] for the
43/// pass-through case. The redactor is `Clone` and `Send + Sync`, so
44/// one instance can be shared across the agent loop for the lifetime
45/// of a run.
46#[derive(Clone)]
47pub struct PayloadRedactor {
48    detector: Arc<dyn PiiDetector>,
49    is_noop: bool,
50}
51
52impl PayloadRedactor {
53    /// Wrap an existing detector.
54    #[must_use]
55    pub fn new(detector: Arc<dyn PiiDetector>) -> Self {
56        Self {
57            detector,
58            is_noop: false,
59        }
60    }
61
62    /// Redactor that performs no masking — produces byte-identical
63    /// JSON output to the raw conversion path.
64    #[must_use]
65    pub fn noop() -> Self {
66        Self {
67            detector: Arc::new(NoopDetector),
68            is_noop: true,
69        }
70    }
71
72    /// Whether this redactor performs no masking.
73    ///
74    /// `true` for [`PayloadRedactor::noop`] / [`Default`]. The
75    /// payload-capture gate uses this to refuse `Inline` payloads from a
76    /// store that attested PII redaction but left the default noop redactor
77    /// in place — the attestation alone is not sufficient evidence that PII
78    /// is actually being masked.
79    #[must_use]
80    pub const fn is_noop(&self) -> bool {
81        self.is_noop
82    }
83
84    /// Convert system instructions, masking PII in the system prompt.
85    ///
86    /// Returns `None` if the system prompt is empty.
87    #[must_use]
88    pub fn convert_system_instructions(&self, request: &ChatRequest) -> Option<Value> {
89        if request.system.is_empty() {
90            return None;
91        }
92        Some(json!([{"text": self.mask_str(&request.system)}]))
93    }
94
95    /// Convert input messages into semconv JSON, masking PII in every
96    /// text and tool-argument leaf.
97    #[must_use]
98    pub fn convert_input_messages(&self, request: &ChatRequest) -> Value {
99        let messages: Vec<Value> = request
100            .messages
101            .iter()
102            .map(|m| self.convert_message(m))
103            .collect();
104        Value::Array(messages)
105    }
106
107    /// Convert a [`ChatResponse`] into semconv output-messages JSON,
108    /// masking PII in every assistant text, thinking text, and
109    /// tool-argument leaf.
110    #[must_use]
111    pub fn convert_output_messages(&self, response: &ChatResponse) -> Value {
112        let parts: Vec<Value> = response
113            .content
114            .iter()
115            .filter_map(|b| self.convert_block(b))
116            .collect();
117        let mut message = json!({
118            "role": "assistant",
119            "content": Value::Array(parts),
120        });
121        if let Some(reason) = response.stop_reason {
122            message["finish_reason"] = json!(finish_reason_str(reason));
123        }
124        json!([message])
125    }
126
127    fn convert_message(&self, message: &Message) -> Value {
128        let role = match message.role {
129            Role::User => determine_user_message_role(message),
130            Role::Assistant => "assistant",
131        };
132        let content = self.convert_content(&message.content);
133        json!({
134            "role": role,
135            "content": content,
136        })
137    }
138
139    fn convert_content(&self, content: &Content) -> Value {
140        match content {
141            Content::Text(text) => json!([{"text": self.mask_str(text)}]),
142            Content::Blocks(blocks) => {
143                let parts: Vec<Value> = blocks
144                    .iter()
145                    .filter_map(|b| self.convert_block(b))
146                    .collect();
147                Value::Array(parts)
148            }
149        }
150    }
151
152    fn convert_block(&self, block: &ContentBlock) -> Option<Value> {
153        // Opaque provider state and redacted-thinking blocks carry nothing
154        // safe to surface. In particular, do not run opaque JSON through the
155        // ordinary PII redactor: ciphertext and future provider-owned fields
156        // must never leave the SDK via telemetry.
157        if matches!(
158            block,
159            ContentBlock::OpaqueReasoning { .. } | ContentBlock::RedactedThinking { .. }
160        ) {
161            return None;
162        }
163
164        match block {
165            ContentBlock::Text { text } => Some(json!({"text": self.mask_str(text)})),
166            ContentBlock::Thinking { thinking, .. } => Some(json!({
167                "type": "reasoning",
168                "text": self.mask_str(thinking),
169            })),
170            ContentBlock::ToolUse {
171                id, name, input, ..
172            } => {
173                let masked_input = self.mask_json(input);
174                Some(json!({
175                    "type": "tool_call",
176                    "id": id,
177                    "name": name,
178                    "arguments": masked_input.to_string(),
179                }))
180            }
181            ContentBlock::ToolResult {
182                tool_use_id,
183                content,
184                is_error,
185            } => {
186                let mut part = json!({
187                    "type": "tool_call_response",
188                    "id": tool_use_id,
189                    "output": self.mask_str(content),
190                });
191                if *is_error == Some(true) {
192                    part["is_error"] = json!(true);
193                }
194                Some(part)
195            }
196            ContentBlock::Image { source } => Some(json!({
197                "type": "blob",
198                "mime_type": source.media_type,
199                "modality": "image",
200                "size": source.data.len(),
201            })),
202            ContentBlock::Document { source } => {
203                let mut part = json!({
204                    "type": "blob",
205                    "mime_type": source.media_type,
206                    "size": source.data.len(),
207                });
208                if source.media_type.starts_with("image/") {
209                    part["modality"] = json!("image");
210                }
211                Some(part)
212            }
213            // Unknown future `#[non_exhaustive]` block kinds are likewise
214            // omitted from the observability payload.
215            _ => None,
216        }
217    }
218
219    /// Mask PII in a plain string.
220    fn mask_str(&self, text: &str) -> String {
221        let spans = self.detector.detect(text);
222        if spans.is_empty() {
223            text.to_owned()
224        } else {
225            mask_spans(text, &spans)
226        }
227    }
228
229    /// Recursively mask every string leaf within a JSON value.
230    fn mask_json(&self, value: &Value) -> Value {
231        match value {
232            Value::String(s) => Value::String(self.mask_str(s)),
233            Value::Array(arr) => Value::Array(arr.iter().map(|v| self.mask_json(v)).collect()),
234            Value::Object(map) => Value::Object(
235                map.iter()
236                    .map(|(k, v)| (k.clone(), self.mask_json(v)))
237                    .collect(),
238            ),
239            Value::Null | Value::Bool(_) | Value::Number(_) => value.clone(),
240        }
241    }
242}
243
244impl Default for PayloadRedactor {
245    fn default() -> Self {
246        Self::noop()
247    }
248}
249
250/// Decide whether a User-role message is actually a `tool` message
251/// (SDK batches tool results as User messages).
252fn determine_user_message_role(message: &Message) -> &'static str {
253    match &message.content {
254        Content::Blocks(blocks) => {
255            let has_tool_result = blocks
256                .iter()
257                .any(|b| matches!(b, ContentBlock::ToolResult { .. }));
258            if has_tool_result { "tool" } else { "user" }
259        }
260        Content::Text(_) => "user",
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use crate::llm::{ChatRequest, ChatResponse, ContentSource, StopReason, Usage};
268    use agent_sdk_foundation::privacy::BaselineDetector;
269
270    fn empty_request(system: &str, messages: Vec<Message>) -> ChatRequest {
271        ChatRequest {
272            system: system.to_owned(),
273            messages,
274            tools: None,
275            max_tokens: 1024,
276            max_tokens_explicit: false,
277            session_id: None,
278            cached_content: None,
279            thinking: None,
280            tool_choice: None,
281            response_format: None,
282            cache: None,
283        }
284    }
285
286    // ── Noop redactor / free functions preserve existing behaviour ──
287
288    #[test]
289    fn empty_system_returns_none() {
290        let request = empty_request("", vec![]);
291        assert!(
292            PayloadRedactor::noop()
293                .convert_system_instructions(&request)
294                .is_none()
295        );
296    }
297
298    #[test]
299    fn system_instructions_wraps_in_text_array() -> anyhow::Result<()> {
300        use anyhow::Context as _;
301        let request = empty_request("You are helpful.", vec![]);
302        let result = PayloadRedactor::noop()
303            .convert_system_instructions(&request)
304            .context("should be Some")?;
305        assert_eq!(result, json!([{"text": "You are helpful."}]));
306        Ok(())
307    }
308
309    #[test]
310    fn user_text_message_converts_correctly() {
311        let msg = Message::user("Hello");
312        let result = PayloadRedactor::noop().convert_message(&msg);
313        assert_eq!(result["role"], "user");
314        assert_eq!(result["content"][0]["text"], "Hello");
315    }
316
317    #[test]
318    fn assistant_text_message_converts_correctly() {
319        let msg = Message::assistant("Hi there");
320        let result = PayloadRedactor::noop().convert_message(&msg);
321        assert_eq!(result["role"], "assistant");
322        assert_eq!(result["content"][0]["text"], "Hi there");
323    }
324
325    #[test]
326    fn tool_result_batch_maps_to_tool_role() {
327        let msg = Message {
328            role: Role::User,
329            content: Content::Blocks(vec![ContentBlock::ToolResult {
330                tool_use_id: "call_1".to_string(),
331                content: "result data".to_string(),
332                is_error: None,
333            }]),
334        };
335        let result = PayloadRedactor::noop().convert_message(&msg);
336        assert_eq!(result["role"], "tool");
337        assert_eq!(result["content"][0]["type"], "tool_call_response");
338        assert_eq!(result["content"][0]["id"], "call_1");
339        assert_eq!(result["content"][0]["output"], "result data");
340    }
341
342    #[test]
343    fn tool_result_with_image_attachment_stays_in_tool_message() {
344        let msg = Message {
345            role: Role::User,
346            content: Content::Blocks(vec![
347                ContentBlock::ToolResult {
348                    tool_use_id: "call_1".to_string(),
349                    content: "screenshot taken".to_string(),
350                    is_error: None,
351                },
352                ContentBlock::Image {
353                    source: ContentSource::new("image/png", "aWdv"),
354                },
355            ]),
356        };
357        let result = PayloadRedactor::noop().convert_message(&msg);
358        assert_eq!(result["role"], "tool");
359        assert_eq!(result["content"][0]["type"], "tool_call_response");
360        assert_eq!(result["content"][1]["type"], "blob");
361        assert_eq!(result["content"][1]["modality"], "image");
362    }
363
364    #[test]
365    fn thinking_block_maps_to_reasoning_part() {
366        let msg = Message {
367            role: Role::Assistant,
368            content: Content::Blocks(vec![
369                ContentBlock::Thinking {
370                    thinking: "Let me think...".to_string(),
371                    signature: None,
372                },
373                ContentBlock::Text {
374                    text: "The answer is 42".to_string(),
375                },
376            ]),
377        };
378        let result = PayloadRedactor::noop().convert_message(&msg);
379        assert_eq!(result["content"][0]["type"], "reasoning");
380        assert_eq!(result["content"][0]["text"], "Let me think...");
381        assert_eq!(result["content"][1]["text"], "The answer is 42");
382    }
383
384    #[test]
385    fn redacted_thinking_is_omitted() {
386        let msg = Message {
387            role: Role::Assistant,
388            content: Content::Blocks(vec![
389                ContentBlock::RedactedThinking {
390                    data: "secret".to_string(),
391                },
392                ContentBlock::Text {
393                    text: "visible".to_string(),
394                },
395            ]),
396        };
397        let result = PayloadRedactor::noop().convert_message(&msg);
398        let content = result["content"].as_array().expect("array");
399        assert_eq!(content.len(), 1);
400        assert_eq!(content[0]["text"], "visible");
401    }
402
403    #[test]
404    fn opaque_reasoning_is_omitted_from_observability_payloads() {
405        let secret = "opaque-secret-that-must-not-leave-the-sdk";
406        let message = Message {
407            role: Role::Assistant,
408            content: Content::Blocks(vec![
409                ContentBlock::OpaqueReasoning {
410                    provider: "test-provider".to_owned(),
411                    data: json!({"encrypted_content": secret}),
412                },
413                ContentBlock::Text {
414                    text: "visible".to_owned(),
415                },
416            ]),
417        };
418        let request = empty_request("", vec![message]);
419        let payload = PayloadRedactor::noop().convert_input_messages(&request);
420
421        assert_eq!(payload[0]["content"].as_array().map(Vec::len), Some(1));
422        assert_eq!(payload[0]["content"][0]["text"], "visible");
423        assert!(!payload.to_string().contains(secret));
424    }
425
426    #[test]
427    fn tool_use_block_maps_to_tool_call_part() {
428        let msg = Message {
429            role: Role::Assistant,
430            content: Content::Blocks(vec![ContentBlock::ToolUse {
431                id: "call_1".to_string(),
432                name: "read".to_string(),
433                input: json!({"path": "/tmp/test.rs"}),
434                thought_signature: None,
435            }]),
436        };
437        let result = PayloadRedactor::noop().convert_message(&msg);
438        assert_eq!(result["content"][0]["type"], "tool_call");
439        assert_eq!(result["content"][0]["id"], "call_1");
440        assert_eq!(result["content"][0]["name"], "read");
441    }
442
443    #[test]
444    fn document_block_maps_to_blob_part() {
445        let msg = Message {
446            role: Role::User,
447            content: Content::Blocks(vec![ContentBlock::Document {
448                source: ContentSource::new("application/pdf", "cGRm"),
449            }]),
450        };
451        let result = PayloadRedactor::noop().convert_message(&msg);
452        assert_eq!(result["content"][0]["type"], "blob");
453        assert_eq!(result["content"][0]["mime_type"], "application/pdf");
454        assert_eq!(result["content"][0]["size"], 4);
455    }
456
457    #[test]
458    fn output_messages_includes_finish_reason() {
459        let response = ChatResponse {
460            id: "resp_1".to_string(),
461            content: vec![ContentBlock::Text {
462                text: "Done".to_string(),
463            }],
464            model: "test-model".to_string(),
465            stop_reason: Some(StopReason::EndTurn),
466            usage: Usage {
467                input_tokens: 10,
468                output_tokens: 5,
469                cached_input_tokens: 0,
470                cache_creation_input_tokens: 0,
471            },
472        };
473        let result = PayloadRedactor::noop().convert_output_messages(&response);
474        let msg = &result[0];
475        assert_eq!(msg["role"], "assistant");
476        assert_eq!(msg["finish_reason"], "stop");
477        assert_eq!(msg["content"][0]["text"], "Done");
478    }
479
480    #[test]
481    fn output_messages_tool_call_finish_reason() {
482        let response = ChatResponse {
483            id: "resp_1".to_string(),
484            content: vec![ContentBlock::ToolUse {
485                id: "c1".to_string(),
486                name: "bash".to_string(),
487                input: json!({"command": "ls"}),
488                thought_signature: None,
489            }],
490            model: "test-model".to_string(),
491            stop_reason: Some(StopReason::ToolUse),
492            usage: Usage {
493                input_tokens: 10,
494                output_tokens: 5,
495                cached_input_tokens: 0,
496                cache_creation_input_tokens: 0,
497            },
498        };
499        let result = PayloadRedactor::noop().convert_output_messages(&response);
500        assert_eq!(result[0]["finish_reason"], "tool_call");
501    }
502
503    #[test]
504    fn tool_result_error_flag_is_preserved() {
505        let msg = Message {
506            role: Role::User,
507            content: Content::Blocks(vec![ContentBlock::ToolResult {
508                tool_use_id: "call_1".to_string(),
509                content: "failed".to_string(),
510                is_error: Some(true),
511            }]),
512        };
513        let result = PayloadRedactor::noop().convert_message(&msg);
514        assert_eq!(result["content"][0]["is_error"], true);
515    }
516
517    #[test]
518    fn input_messages_preserves_order() {
519        let request = empty_request(
520            "",
521            vec![
522                Message::user("first"),
523                Message::assistant("second"),
524                Message::user("third"),
525            ],
526        );
527        let result = PayloadRedactor::noop().convert_input_messages(&request);
528        let arr = result.as_array().expect("array");
529        assert_eq!(arr.len(), 3);
530        assert_eq!(arr[0]["role"], "user");
531        assert_eq!(arr[1]["role"], "assistant");
532        assert_eq!(arr[2]["role"], "user");
533    }
534
535    // ── Baseline redactor masks PII across every covered surface ──
536
537    fn baseline_redactor() -> PayloadRedactor {
538        PayloadRedactor::new(Arc::new(
539            BaselineDetector::new().expect("baseline compiles"),
540        ))
541    }
542
543    #[test]
544    fn redacts_email_in_system_prompt() {
545        let request = empty_request("Contact support at ops@example.com.", vec![]);
546        let result = baseline_redactor()
547            .convert_system_instructions(&request)
548            .expect("some");
549        let text = result[0]["text"].as_str().expect("text string");
550        assert!(
551            text.contains("[REDACTED:email]"),
552            "system prompt not redacted: {text}"
553        );
554        assert!(!text.contains("ops@example.com"));
555    }
556
557    #[test]
558    fn redacts_cpf_in_user_text() {
559        let request = empty_request("", vec![Message::user("meu CPF é 111.444.777-35")]);
560        let result = baseline_redactor().convert_input_messages(&request);
561        let text = result[0]["content"][0]["text"].as_str().expect("text");
562        assert!(text.contains("[REDACTED:cpf]"), "user text: {text}");
563        assert!(!text.contains("111.444.777-35"));
564    }
565
566    #[test]
567    fn redacts_pan_in_tool_result_output() {
568        let msg = Message {
569            role: Role::User,
570            content: Content::Blocks(vec![ContentBlock::ToolResult {
571                tool_use_id: "c1".to_string(),
572                content: "charged card 4111 1111 1111 1111 successfully".to_string(),
573                is_error: None,
574            }]),
575        };
576        let result = baseline_redactor().convert_message(&msg);
577        let output = result["content"][0]["output"].as_str().expect("output");
578        assert!(
579            output.contains("[REDACTED:credit_card]"),
580            "tool output: {output}"
581        );
582        assert!(!output.contains("4111 1111 1111 1111"));
583    }
584
585    #[test]
586    fn redacts_strings_inside_tool_call_arguments_json() {
587        let msg = Message {
588            role: Role::Assistant,
589            content: Content::Blocks(vec![ContentBlock::ToolUse {
590                id: "c1".to_string(),
591                name: "send_pix".to_string(),
592                input: json!({
593                    "chave_pix": "ana@example.com",
594                    "amount_brl": 100,
595                    "metadata": {
596                        "recipient_cpf": "111.444.777-35",
597                        "note": "salário"
598                    }
599                }),
600                thought_signature: None,
601            }]),
602        };
603        let result = baseline_redactor().convert_message(&msg);
604        let args = result["content"][0]["arguments"].as_str().expect("args");
605        assert!(args.contains("[REDACTED:email]"), "args: {args}");
606        assert!(args.contains("[REDACTED:cpf]"), "args: {args}");
607        assert!(!args.contains("ana@example.com"));
608        assert!(!args.contains("111.444.777-35"));
609        // Non-string leaves are preserved as-is.
610        assert!(args.contains("100"));
611    }
612
613    #[test]
614    fn redacts_secret_in_assistant_output() {
615        let response = ChatResponse {
616            id: "r1".to_string(),
617            content: vec![ContentBlock::Text {
618                text: "here is the key sk-abcdefghijklmnopqrstuv for ci".to_string(),
619            }],
620            model: "m".to_string(),
621            stop_reason: Some(StopReason::EndTurn),
622            usage: Usage {
623                input_tokens: 0,
624                output_tokens: 0,
625                cached_input_tokens: 0,
626                cache_creation_input_tokens: 0,
627            },
628        };
629        let result = baseline_redactor().convert_output_messages(&response);
630        let text = result[0]["content"][0]["text"].as_str().expect("text");
631        assert!(text.contains("[REDACTED:secret]"), "output: {text}");
632        assert!(!text.contains("sk-abcdefghijklmnopqrstuv"));
633    }
634
635    #[test]
636    fn redacts_pii_in_thinking_text() {
637        let msg = Message {
638            role: Role::Assistant,
639            content: Content::Blocks(vec![ContentBlock::Thinking {
640                thinking: "User CPF is 111.444.777-35 — I should confirm before sending."
641                    .to_string(),
642                signature: None,
643            }]),
644        };
645        let result = baseline_redactor().convert_message(&msg);
646        let text = result["content"][0]["text"].as_str().expect("text");
647        assert!(text.contains("[REDACTED:cpf]"), "thinking: {text}");
648    }
649
650    #[test]
651    fn mask_json_preserves_non_string_leaves() {
652        let input = json!({
653            "amount": 42.5,
654            "active": true,
655            "items": null,
656            "email": "ana@example.com"
657        });
658        let redacted = baseline_redactor().mask_json(&input);
659        assert_eq!(redacted["amount"], json!(42.5));
660        assert_eq!(redacted["active"], json!(true));
661        assert_eq!(redacted["items"], json!(null));
662        assert!(
663            redacted["email"]
664                .as_str()
665                .expect("email")
666                .contains("[REDACTED:email]")
667        );
668    }
669
670    #[test]
671    fn is_noop_reflects_constructor() {
672        assert!(PayloadRedactor::noop().is_noop());
673        assert!(PayloadRedactor::default().is_noop());
674        assert!(
675            !baseline_redactor().is_noop(),
676            "a detector-backed redactor must not report as noop"
677        );
678    }
679}