1use 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#[derive(Clone)]
47pub struct PayloadRedactor {
48 detector: Arc<dyn PiiDetector>,
49 is_noop: bool,
50}
51
52impl PayloadRedactor {
53 #[must_use]
55 pub fn new(detector: Arc<dyn PiiDetector>) -> Self {
56 Self {
57 detector,
58 is_noop: false,
59 }
60 }
61
62 #[must_use]
65 pub fn noop() -> Self {
66 Self {
67 detector: Arc::new(NoopDetector),
68 is_noop: true,
69 }
70 }
71
72 #[must_use]
80 pub const fn is_noop(&self) -> bool {
81 self.is_noop
82 }
83
84 #[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 #[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 #[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 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 _ => None,
216 }
217 }
218
219 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 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
250fn 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 #[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 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 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}