1use serde::{Deserialize, Serialize};
22
23#[cfg(feature = "openapi")]
24use utoipa::ToSchema;
25
26pub const MAX_CHECKS: usize = 64;
28pub const MAX_ENTRIES_PER_CHECK: usize = 64;
30pub const MAX_ENTRY_LEN: usize = 512;
32pub const MAX_REPLACEMENT_LEN: usize = 2_000;
34pub const MAX_CHECK_ID_LEN: usize = 64;
36pub const MAX_JUDGE_PROMPT_LEN: usize = 4_000;
38pub const MAX_MCP_REF_LEN: usize = 128;
40const REGEX_SIZE_LIMIT: usize = 1 << 20;
43const MAX_MATCH_SNIPPET: usize = 200;
46
47pub const DEFAULT_OUTPUT_REPLACEMENT: &str = "[Response withheld by a guardrail.]";
49pub const DEFAULT_TOOL_OUTPUT_REPLACEMENT: &str = "[Tool output withheld by a guardrail.]";
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
56#[cfg_attr(feature = "openapi", derive(ToSchema))]
57#[serde(rename_all = "snake_case")]
58pub enum GuardrailMode {
59 #[default]
60 Active,
61 Advisory,
62}
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
66#[cfg_attr(feature = "openapi", derive(ToSchema))]
67#[cfg_attr(feature = "openapi", schema(example = "output"))]
68#[serde(rename_all = "snake_case")]
69pub enum GuardrailStage {
70 Output,
73 ToolUse,
75 ToolOutput,
77}
78
79impl GuardrailStage {
80 pub fn as_str(&self) -> &'static str {
81 match self {
82 GuardrailStage::Output => "output",
83 GuardrailStage::ToolUse => "tool_use",
84 GuardrailStage::ToolOutput => "tool_output",
85 }
86 }
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
92#[cfg_attr(feature = "openapi", derive(ToSchema))]
93#[serde(rename_all = "snake_case")]
94pub enum GuardrailOnFail {
95 #[default]
96 Block,
97 Log,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
107#[serde(tag = "type", rename_all = "snake_case")]
108pub enum GuardrailRule {
109 Regex { patterns: Vec<String> },
111 Blocklist {
113 words: Vec<String>,
114 #[serde(default)]
115 case_sensitive: bool,
116 },
117 ToolPattern { tools: Vec<String> },
120 LlmJudge { prompt: String },
127 Mcp { server: String, tool: String },
141 Moderation {
153 #[serde(default)]
154 categories: Vec<String>,
155 #[serde(default = "default_moderation_threshold")]
156 threshold: u8,
157 },
158}
159
160pub fn default_moderation_threshold() -> u8 {
162 50
163}
164
165pub const DEFAULT_MODERATION_CATEGORIES: &[&str] = &[
167 "hate",
168 "harassment",
169 "self_harm",
170 "sexual",
171 "violence",
172 "illicit",
173];
174
175impl GuardrailRule {
176 pub fn rule_type(&self) -> &'static str {
177 match self {
178 GuardrailRule::Regex { .. } => "regex",
179 GuardrailRule::Blocklist { .. } => "blocklist",
180 GuardrailRule::ToolPattern { .. } => "tool_pattern",
181 GuardrailRule::LlmJudge { .. } => "llm_judge",
182 GuardrailRule::Mcp { .. } => "mcp",
183 GuardrailRule::Moderation { .. } => "moderation",
184 }
185 }
186}
187
188#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
190pub struct GuardrailCheck {
191 #[serde(default, skip_serializing_if = "Option::is_none")]
193 pub id: Option<String>,
194 pub stage: GuardrailStage,
195 #[serde(default)]
196 pub on_fail: GuardrailOnFail,
197 #[serde(default, skip_serializing_if = "Option::is_none")]
200 pub replacement: Option<String>,
201 #[serde(flatten)]
202 pub rule: GuardrailRule,
203}
204
205#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
208pub struct GuardrailsConfig {
209 #[serde(default)]
210 pub mode: GuardrailMode,
211 #[serde(default)]
212 pub checks: Vec<GuardrailCheck>,
213}
214
215impl GuardrailsConfig {
216 pub fn from_value(value: &serde_json::Value) -> Result<Self, String> {
218 if value.is_null() {
219 return Ok(Self::default());
220 }
221 serde_json::from_value(value.clone()).map_err(|e| format!("invalid guardrails config: {e}"))
222 }
223
224 pub fn compile(&self) -> Result<CompiledGuardrails, String> {
227 if self.checks.len() > MAX_CHECKS {
228 return Err(format!(
229 "too many checks: {} (max {MAX_CHECKS})",
230 self.checks.len()
231 ));
232 }
233 let mut compiled = Vec::with_capacity(self.checks.len());
234 let mut judge_checks = Vec::new();
235 let mut mcp_checks = Vec::new();
236 let mut moderation_checks = Vec::new();
237 for (index, check) in self.checks.iter().enumerate() {
238 match &check.rule {
239 GuardrailRule::LlmJudge { prompt } => {
240 judge_checks.push(compile_judge_check(index, check, prompt)?);
241 }
242 GuardrailRule::Mcp { server, tool } => {
243 mcp_checks.push(compile_mcp_check(index, check, server, tool)?);
244 }
245 GuardrailRule::Moderation {
246 categories,
247 threshold,
248 } => {
249 moderation_checks.push(compile_moderation_check(
250 index, check, categories, *threshold,
251 )?);
252 }
253 _ => compiled.push(compile_check(index, check)?),
254 }
255 }
256 Ok(CompiledGuardrails {
257 mode: self.mode,
258 checks: compiled,
259 judge_checks,
260 mcp_checks,
261 moderation_checks,
262 })
263 }
264}
265
266#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
268#[cfg_attr(feature = "openapi", derive(ToSchema))]
269#[cfg_attr(feature = "openapi", schema(example = "block"))]
270#[serde(rename_all = "snake_case")]
271pub enum GuardrailAction {
272 Block,
273 Log,
274}
275
276#[derive(Debug, Clone, PartialEq, Eq)]
278pub struct GuardrailHit {
279 pub check_index: usize,
281 pub check_label: String,
283 pub stage: GuardrailStage,
284 pub rule_type: &'static str,
285 pub action: GuardrailAction,
287 pub reason_code: String,
289 pub replacement: Option<String>,
291 pub matched: Option<String>,
294}
295
296#[derive(Debug)]
299pub struct CompiledJudgeCheck {
300 pub index: usize,
301 pub label: String,
302 pub stage: GuardrailStage,
303 pub on_fail: GuardrailOnFail,
304 pub replacement: Option<String>,
305 pub prompt: String,
307}
308
309#[derive(Debug)]
313pub struct CompiledMcpCheck {
314 pub index: usize,
315 pub label: String,
316 pub stage: GuardrailStage,
317 pub on_fail: GuardrailOnFail,
318 pub replacement: Option<String>,
319 pub server: String,
321 pub tool: String,
323}
324
325#[derive(Debug)]
329pub struct CompiledModerationCheck {
330 pub index: usize,
331 pub label: String,
332 pub stage: GuardrailStage,
333 pub on_fail: GuardrailOnFail,
334 pub replacement: Option<String>,
335 pub categories: Vec<String>,
338 pub threshold: u8,
341}
342
343#[derive(Debug)]
345pub struct CompiledGuardrails {
346 mode: GuardrailMode,
347 checks: Vec<CompiledCheck>,
348 judge_checks: Vec<CompiledJudgeCheck>,
351 mcp_checks: Vec<CompiledMcpCheck>,
354 moderation_checks: Vec<CompiledModerationCheck>,
357}
358
359#[derive(Debug)]
360struct CompiledCheck {
361 index: usize,
362 label: String,
363 stage: GuardrailStage,
364 on_fail: GuardrailOnFail,
365 replacement: Option<String>,
366 rule_type: &'static str,
367 matcher: CompiledRule,
368}
369
370#[derive(Debug)]
371enum CompiledRule {
372 Regex(Vec<regex::Regex>),
373 Blocklist {
374 words: Vec<String>,
376 case_sensitive: bool,
377 },
378 ToolPattern(Vec<String>),
379}
380
381impl CompiledGuardrails {
382 pub fn mode(&self) -> GuardrailMode {
383 self.mode
384 }
385
386 pub fn has_stage(&self, stage: GuardrailStage) -> bool {
389 self.checks.iter().any(|c| c.stage == stage)
390 || self.judge_checks.iter().any(|c| c.stage == stage)
391 || self.mcp_checks.iter().any(|c| c.stage == stage)
392 || self.moderation_checks.iter().any(|c| c.stage == stage)
393 }
394
395 pub fn moderation_checks_for_stage(
399 &self,
400 stage: GuardrailStage,
401 ) -> impl Iterator<Item = &CompiledModerationCheck> {
402 self.moderation_checks
403 .iter()
404 .filter(move |c| c.stage == stage)
405 }
406
407 pub fn judge_checks_for_stage(
410 &self,
411 stage: GuardrailStage,
412 ) -> impl Iterator<Item = &CompiledJudgeCheck> {
413 self.judge_checks.iter().filter(move |c| c.stage == stage)
414 }
415
416 pub fn mcp_checks_for_stage(
419 &self,
420 stage: GuardrailStage,
421 ) -> impl Iterator<Item = &CompiledMcpCheck> {
422 self.mcp_checks.iter().filter(move |c| c.stage == stage)
423 }
424
425 pub fn async_action(&self, on_fail: GuardrailOnFail) -> GuardrailAction {
428 match (self.mode, on_fail) {
429 (GuardrailMode::Advisory, _) | (_, GuardrailOnFail::Log) => GuardrailAction::Log,
430 (GuardrailMode::Active, GuardrailOnFail::Block) => GuardrailAction::Block,
431 }
432 }
433
434 pub fn judge_action(&self, on_fail: GuardrailOnFail) -> GuardrailAction {
437 self.async_action(on_fail)
438 }
439
440 pub fn evaluate(
446 &self,
447 stage: GuardrailStage,
448 text: &str,
449 tool_name: Option<&str>,
450 skip: &dyn Fn(usize) -> bool,
451 ) -> Vec<GuardrailHit> {
452 let lowercased: std::cell::OnceCell<String> = std::cell::OnceCell::new();
453 let mut hits = Vec::new();
454 for check in self.checks.iter() {
455 if check.stage != stage || skip(check.index) {
456 continue;
457 }
458 let matched = match &check.matcher {
459 CompiledRule::Regex(patterns) => patterns
460 .iter()
461 .find_map(|re| re.find(text).map(|m| snippet(m.as_str()))),
462 CompiledRule::Blocklist {
463 words,
464 case_sensitive,
465 } => {
466 let haystack: &str = if *case_sensitive {
467 text
468 } else {
469 lowercased.get_or_init(|| text.to_lowercase())
470 };
471 words
472 .iter()
473 .find(|w| haystack.contains(w.as_str()))
474 .map(|w| snippet(w))
475 }
476 CompiledRule::ToolPattern(patterns) => tool_name.and_then(|name| {
477 patterns
478 .iter()
479 .find(|p| wildcard_match(p, name))
480 .map(|_| snippet(name))
481 }),
482 };
483 if matched.is_some() {
484 let action = match (self.mode, check.on_fail) {
485 (GuardrailMode::Advisory, _) | (_, GuardrailOnFail::Log) => {
486 GuardrailAction::Log
487 }
488 (GuardrailMode::Active, GuardrailOnFail::Block) => GuardrailAction::Block,
489 };
490 hits.push(GuardrailHit {
491 check_index: check.index,
492 check_label: check.label.clone(),
493 stage: check.stage,
494 rule_type: check.rule_type,
495 action,
496 reason_code: format!("guardrail.{}", check.rule_type),
497 replacement: check.replacement.clone(),
498 matched,
499 });
500 }
501 }
502 hits
503 }
504}
505
506fn compile_check(index: usize, check: &GuardrailCheck) -> Result<CompiledCheck, String> {
507 let label = match &check.id {
508 Some(id) => {
509 if id.is_empty() || id.chars().count() > MAX_CHECK_ID_LEN {
510 return Err(format!(
511 "check #{index}: id must be 1..={MAX_CHECK_ID_LEN} characters"
512 ));
513 }
514 id.clone()
515 }
516 None => format!("{}#{}", check.rule.rule_type(), index),
517 };
518 if let Some(replacement) = &check.replacement
519 && replacement.len() > MAX_REPLACEMENT_LEN
520 {
521 return Err(format!(
522 "check '{label}': replacement exceeds {MAX_REPLACEMENT_LEN} bytes"
523 ));
524 }
525 let matcher = match &check.rule {
526 GuardrailRule::Regex { patterns } => {
527 validate_entries(&label, "patterns", patterns)?;
528 let mut compiled = Vec::with_capacity(patterns.len());
529 for pattern in patterns {
530 let re = regex::RegexBuilder::new(pattern)
531 .size_limit(REGEX_SIZE_LIMIT)
532 .build()
533 .map_err(|e| format!("check '{label}': invalid regex '{pattern}': {e}"))?;
534 compiled.push(re);
535 }
536 CompiledRule::Regex(compiled)
537 }
538 GuardrailRule::Blocklist {
539 words,
540 case_sensitive,
541 } => {
542 validate_entries(&label, "words", words)?;
543 let words = if *case_sensitive {
544 words.clone()
545 } else {
546 words.iter().map(|w| w.to_lowercase()).collect()
547 };
548 CompiledRule::Blocklist {
549 words,
550 case_sensitive: *case_sensitive,
551 }
552 }
553 GuardrailRule::ToolPattern { tools } => {
554 if check.stage != GuardrailStage::ToolUse {
555 return Err(format!(
556 "check '{label}': tool_pattern is only valid for the tool_use stage"
557 ));
558 }
559 validate_entries(&label, "tools", tools)?;
560 CompiledRule::ToolPattern(tools.clone())
561 }
562 GuardrailRule::LlmJudge { .. } => {
563 unreachable!(
564 "llm_judge checks are routed to compile_judge_check before compile_check is called"
565 )
566 }
567 GuardrailRule::Mcp { .. } => {
568 unreachable!(
569 "mcp checks are routed to compile_mcp_check before compile_check is called"
570 )
571 }
572 GuardrailRule::Moderation { .. } => {
573 unreachable!(
574 "moderation checks are routed to compile_moderation_check before compile_check is called"
575 )
576 }
577 };
578 Ok(CompiledCheck {
579 index,
580 label,
581 stage: check.stage,
582 on_fail: check.on_fail,
583 replacement: check.replacement.clone(),
584 rule_type: check.rule.rule_type(),
585 matcher,
586 })
587}
588
589fn compile_judge_check(
590 index: usize,
591 check: &GuardrailCheck,
592 prompt: &str,
593) -> Result<CompiledJudgeCheck, String> {
594 let label = match &check.id {
595 Some(id) => {
596 if id.is_empty() || id.chars().count() > MAX_CHECK_ID_LEN {
597 return Err(format!(
598 "check #{index}: id must be 1..={MAX_CHECK_ID_LEN} characters"
599 ));
600 }
601 id.clone()
602 }
603 None => format!("llm_judge#{index}"),
604 };
605 if prompt.is_empty() {
606 return Err(format!(
607 "check '{label}': llm_judge prompt must not be empty"
608 ));
609 }
610 if prompt.len() > MAX_JUDGE_PROMPT_LEN {
611 return Err(format!(
612 "check '{label}': llm_judge prompt exceeds {MAX_JUDGE_PROMPT_LEN} bytes"
613 ));
614 }
615 match check.stage {
616 GuardrailStage::ToolUse | GuardrailStage::ToolOutput => {}
617 GuardrailStage::Output => {
618 return Err(format!(
619 "check '{label}': llm_judge is not supported on the 'output' stage in this phase; \
620 use 'tool_use' or 'tool_output'"
621 ));
622 }
623 }
624 if let Some(replacement) = &check.replacement
625 && replacement.len() > MAX_REPLACEMENT_LEN
626 {
627 return Err(format!(
628 "check '{label}': replacement exceeds {MAX_REPLACEMENT_LEN} bytes"
629 ));
630 }
631 Ok(CompiledJudgeCheck {
632 index,
633 label,
634 stage: check.stage,
635 on_fail: check.on_fail,
636 replacement: check.replacement.clone(),
637 prompt: prompt.to_string(),
638 })
639}
640
641fn compile_mcp_check(
642 index: usize,
643 check: &GuardrailCheck,
644 server: &str,
645 tool: &str,
646) -> Result<CompiledMcpCheck, String> {
647 let label = match &check.id {
648 Some(id) => {
649 if id.is_empty() || id.chars().count() > MAX_CHECK_ID_LEN {
650 return Err(format!(
651 "check #{index}: id must be 1..={MAX_CHECK_ID_LEN} characters"
652 ));
653 }
654 id.clone()
655 }
656 None => format!("mcp#{index}"),
657 };
658 match check.stage {
661 GuardrailStage::ToolUse | GuardrailStage::ToolOutput => {}
662 GuardrailStage::Output => {
663 return Err(format!(
664 "check '{label}': mcp is not supported on the 'output' stage in this phase; \
665 use 'tool_use' or 'tool_output'"
666 ));
667 }
668 }
669 for (field, value) in [("server", server), ("tool", tool)] {
670 if value.is_empty() {
671 return Err(format!("check '{label}': mcp {field} must not be empty"));
672 }
673 if value.len() > MAX_MCP_REF_LEN {
674 return Err(format!(
675 "check '{label}': mcp {field} exceeds {MAX_MCP_REF_LEN} bytes"
676 ));
677 }
678 }
679 if let Some(replacement) = &check.replacement
680 && replacement.len() > MAX_REPLACEMENT_LEN
681 {
682 return Err(format!(
683 "check '{label}': replacement exceeds {MAX_REPLACEMENT_LEN} bytes"
684 ));
685 }
686 Ok(CompiledMcpCheck {
687 index,
688 label,
689 stage: check.stage,
690 on_fail: check.on_fail,
691 replacement: check.replacement.clone(),
692 server: server.to_string(),
693 tool: tool.to_string(),
694 })
695}
696
697fn compile_moderation_check(
698 index: usize,
699 check: &GuardrailCheck,
700 categories: &[String],
701 threshold: u8,
702) -> Result<CompiledModerationCheck, String> {
703 let label = match &check.id {
704 Some(id) => {
705 if id.is_empty() || id.chars().count() > MAX_CHECK_ID_LEN {
706 return Err(format!(
707 "check #{index}: id must be 1..={MAX_CHECK_ID_LEN} characters"
708 ));
709 }
710 id.clone()
711 }
712 None => format!("moderation#{index}"),
713 };
714 match check.stage {
717 GuardrailStage::Output => {}
718 GuardrailStage::ToolUse | GuardrailStage::ToolOutput => {
719 return Err(format!(
720 "check '{label}': moderation is only supported on the 'output' stage"
721 ));
722 }
723 }
724 if threshold > 100 {
725 return Err(format!(
726 "check '{label}': moderation threshold must be 0..=100 (got {threshold})"
727 ));
728 }
729 let categories = if categories.is_empty() {
732 DEFAULT_MODERATION_CATEGORIES
733 .iter()
734 .map(|c| (*c).to_string())
735 .collect()
736 } else {
737 validate_entries(&label, "categories", categories)?;
738 categories.to_vec()
739 };
740 if let Some(replacement) = &check.replacement
741 && replacement.len() > MAX_REPLACEMENT_LEN
742 {
743 return Err(format!(
744 "check '{label}': replacement exceeds {MAX_REPLACEMENT_LEN} bytes"
745 ));
746 }
747 Ok(CompiledModerationCheck {
748 index,
749 label,
750 stage: check.stage,
751 on_fail: check.on_fail,
752 replacement: check.replacement.clone(),
753 categories,
754 threshold,
755 })
756}
757
758fn validate_entries(label: &str, field: &str, entries: &[String]) -> Result<(), String> {
759 if entries.is_empty() {
760 return Err(format!("check '{label}': {field} must not be empty"));
761 }
762 if entries.len() > MAX_ENTRIES_PER_CHECK {
763 return Err(format!(
764 "check '{label}': too many {field}: {} (max {MAX_ENTRIES_PER_CHECK})",
765 entries.len()
766 ));
767 }
768 for entry in entries {
769 if entry.is_empty() {
770 return Err(format!(
771 "check '{label}': {field} entries must not be empty"
772 ));
773 }
774 if entry.len() > MAX_ENTRY_LEN {
775 return Err(format!(
776 "check '{label}': {field} entry exceeds {MAX_ENTRY_LEN} bytes"
777 ));
778 }
779 }
780 Ok(())
781}
782
783fn snippet(s: &str) -> String {
784 let mut end = MAX_MATCH_SNIPPET.min(s.len());
785 while end > 0 && !s.is_char_boundary(end) {
786 end -= 1;
787 }
788 s[..end].to_string()
789}
790
791pub fn wildcard_match(pattern: &str, name: &str) -> bool {
794 let segments: Vec<&str> = pattern.split('*').collect();
795 if segments.len() == 1 {
796 return pattern == name;
797 }
798 let mut rest = name;
799 for (i, seg) in segments.iter().enumerate() {
800 if seg.is_empty() {
801 continue;
802 }
803 if i == 0 {
804 match rest.strip_prefix(seg) {
805 Some(r) => rest = r,
806 None => return false,
807 }
808 } else if i == segments.len() - 1 {
809 return rest.ends_with(seg);
810 } else {
811 match rest.find(seg) {
812 Some(pos) => rest = &rest[pos + seg.len()..],
813 None => return false,
814 }
815 }
816 }
817 true
819}
820
821#[cfg(test)]
822mod tests {
823 use super::*;
824 use serde_json::json;
825
826 fn no_skip() -> impl Fn(usize) -> bool {
827 |_| false
828 }
829
830 fn compile(value: serde_json::Value) -> Result<CompiledGuardrails, String> {
831 GuardrailsConfig::from_value(&value)?.compile()
832 }
833
834 #[test]
835 fn parses_and_compiles_minimal_config() {
836 let compiled = compile(json!({
837 "checks": [
838 {"stage": "output", "type": "blocklist", "words": ["forbidden"]},
839 ]
840 }))
841 .expect("compiles");
842 assert_eq!(compiled.mode(), GuardrailMode::Active);
843 assert!(compiled.has_stage(GuardrailStage::Output));
844 assert!(!compiled.has_stage(GuardrailStage::ToolUse));
845 }
846
847 #[test]
848 fn empty_or_null_config_compiles_to_no_checks() {
849 let compiled = compile(json!({})).expect("compiles");
850 assert!(!compiled.has_stage(GuardrailStage::Output));
851 let compiled = GuardrailsConfig::from_value(&serde_json::Value::Null)
852 .unwrap()
853 .compile()
854 .unwrap();
855 assert!(!compiled.has_stage(GuardrailStage::Output));
856 }
857
858 #[test]
859 fn blocklist_matches_case_insensitive_by_default() {
860 let compiled = compile(json!({
861 "checks": [
862 {"stage": "output", "type": "blocklist", "words": ["Secret Word"]},
863 ]
864 }))
865 .unwrap();
866 let hits = compiled.evaluate(
867 GuardrailStage::Output,
868 "this contains a SECRET word inside",
869 None,
870 &no_skip(),
871 );
872 assert_eq!(hits.len(), 1);
873 assert_eq!(hits[0].action, GuardrailAction::Block);
874 assert_eq!(hits[0].reason_code, "guardrail.blocklist");
875 assert_eq!(hits[0].matched.as_deref(), Some("secret word"));
876 }
877
878 #[test]
879 fn blocklist_case_sensitive_only_matches_exact_case() {
880 let compiled = compile(json!({
881 "checks": [
882 {"stage": "output", "type": "blocklist", "words": ["Secret"], "case_sensitive": true},
883 ]
884 }))
885 .unwrap();
886 assert!(
887 compiled
888 .evaluate(GuardrailStage::Output, "a secret here", None, &no_skip())
889 .is_empty()
890 );
891 assert_eq!(
892 compiled
893 .evaluate(GuardrailStage::Output, "a Secret here", None, &no_skip())
894 .len(),
895 1
896 );
897 }
898
899 #[test]
900 fn regex_matches_and_reports_pattern_source() {
901 let compiled = compile(json!({
902 "checks": [
903 {"id": "ssn", "stage": "output", "type": "regex",
904 "patterns": ["\\b\\d{3}-\\d{2}-\\d{4}\\b"]},
905 ]
906 }))
907 .unwrap();
908 let hits = compiled.evaluate(
909 GuardrailStage::Output,
910 "my ssn is 123-45-6789 ok",
911 None,
912 &no_skip(),
913 );
914 assert_eq!(hits.len(), 1);
915 assert_eq!(hits[0].check_label, "ssn");
916 assert_eq!(hits[0].rule_type, "regex");
917 }
918
919 #[test]
920 fn invalid_regex_fails_compile_with_check_label() {
921 let err = compile(json!({
922 "checks": [
923 {"id": "bad", "stage": "output", "type": "regex", "patterns": ["("]},
924 ]
925 }))
926 .unwrap_err();
927 assert!(err.contains("check 'bad'"), "{err}");
928 }
929
930 #[test]
931 fn tool_pattern_matches_wildcards_on_tool_use_stage() {
932 let compiled = compile(json!({
933 "checks": [
934 {"stage": "tool_use", "type": "tool_pattern", "tools": ["mcp_*", "bash*"]},
935 ]
936 }))
937 .unwrap();
938 let hits = compiled.evaluate(
939 GuardrailStage::ToolUse,
940 "{\"cmd\":\"ls\"}",
941 Some("mcp_github__create_issue"),
942 &no_skip(),
943 );
944 assert_eq!(hits.len(), 1);
945 assert_eq!(hits[0].matched.as_deref(), Some("mcp_github__create_issue"));
946 assert!(
947 compiled
948 .evaluate(GuardrailStage::ToolUse, "{}", Some("read_file"), &no_skip())
949 .is_empty()
950 );
951 }
952
953 #[test]
954 fn tool_pattern_rejected_outside_tool_use_stage() {
955 let err = compile(json!({
956 "checks": [
957 {"stage": "output", "type": "tool_pattern", "tools": ["bash*"]},
958 ]
959 }))
960 .unwrap_err();
961 assert!(err.contains("only valid for the tool_use stage"), "{err}");
962 }
963
964 #[test]
965 fn advisory_mode_downgrades_block_to_log() {
966 let compiled = compile(json!({
967 "mode": "advisory",
968 "checks": [
969 {"stage": "output", "type": "blocklist", "words": ["x"], "on_fail": "block"},
970 ]
971 }))
972 .unwrap();
973 let hits = compiled.evaluate(GuardrailStage::Output, "x", None, &no_skip());
974 assert_eq!(hits[0].action, GuardrailAction::Log);
975 }
976
977 #[test]
978 fn on_fail_log_yields_log_action_in_active_mode() {
979 let compiled = compile(json!({
980 "checks": [
981 {"stage": "output", "type": "blocklist", "words": ["x"], "on_fail": "log"},
982 ]
983 }))
984 .unwrap();
985 let hits = compiled.evaluate(GuardrailStage::Output, "x", None, &no_skip());
986 assert_eq!(hits[0].action, GuardrailAction::Log);
987 }
988
989 #[test]
990 fn skip_suppresses_checks_by_index() {
991 let compiled = compile(json!({
992 "checks": [
993 {"stage": "output", "type": "blocklist", "words": ["x"]},
994 {"stage": "output", "type": "blocklist", "words": ["y"]},
995 ]
996 }))
997 .unwrap();
998 let hits = compiled.evaluate(GuardrailStage::Output, "x y", None, &|i| i == 0);
999 assert_eq!(hits.len(), 1);
1000 assert_eq!(hits[0].check_index, 1);
1001 }
1002
1003 #[test]
1004 fn enforces_limits() {
1005 let too_many_checks: Vec<_> = (0..=MAX_CHECKS)
1006 .map(|_| json!({"stage": "output", "type": "blocklist", "words": ["x"]}))
1007 .collect();
1008 assert!(
1009 compile(json!({"checks": too_many_checks}))
1010 .unwrap_err()
1011 .contains("too many checks")
1012 );
1013
1014 let long_entry = "a".repeat(MAX_ENTRY_LEN + 1);
1015 assert!(
1016 compile(json!({
1017 "checks": [{"stage": "output", "type": "blocklist", "words": [long_entry]}]
1018 }))
1019 .unwrap_err()
1020 .contains("exceeds")
1021 );
1022
1023 assert!(
1024 compile(json!({
1025 "checks": [{"stage": "output", "type": "blocklist", "words": []}]
1026 }))
1027 .unwrap_err()
1028 .contains("must not be empty")
1029 );
1030 }
1031
1032 #[test]
1033 fn unknown_fields_are_rejected_gracefully_by_value_parse() {
1034 let err = GuardrailsConfig::from_value(&json!({"checks": "nope"})).unwrap_err();
1037 assert!(err.contains("invalid guardrails config"), "{err}");
1038 }
1039
1040 #[test]
1041 fn wildcard_match_covers_anchors_and_inner_stars() {
1042 assert!(wildcard_match("bash*", "bashkit_exec"));
1043 assert!(wildcard_match("*_file", "read_file"));
1044 assert!(wildcard_match("mcp_*__delete_*", "mcp_github__delete_repo"));
1045 assert!(wildcard_match("*", "anything"));
1046 assert!(wildcard_match("exact", "exact"));
1047 assert!(!wildcard_match("exact", "exact_no"));
1048 assert!(!wildcard_match("bash*", "zsh"));
1049 assert!(!wildcard_match(
1050 "mcp_*__delete_*",
1051 "mcp_github__create_repo"
1052 ));
1053 }
1054
1055 #[test]
1056 fn config_roundtrips_serde() {
1057 let cfg = GuardrailsConfig {
1058 mode: GuardrailMode::Advisory,
1059 checks: vec![GuardrailCheck {
1060 id: Some("c1".into()),
1061 stage: GuardrailStage::ToolUse,
1062 on_fail: GuardrailOnFail::Log,
1063 replacement: None,
1064 rule: GuardrailRule::ToolPattern {
1065 tools: vec!["bash*".into()],
1066 },
1067 }],
1068 };
1069 let value = serde_json::to_value(&cfg).unwrap();
1070 assert_eq!(value["checks"][0]["type"], "tool_pattern");
1071 assert_eq!(value["checks"][0]["stage"], "tool_use");
1072 let back = GuardrailsConfig::from_value(&value).unwrap();
1073 assert_eq!(back, cfg);
1074 }
1075
1076 #[test]
1079 fn llm_judge_compiles_for_tool_stages() {
1080 let compiled = compile(json!({
1081 "checks": [
1082 {"stage": "tool_use", "type": "llm_judge", "prompt": "Block requests to delete data."},
1083 {"id": "tj2", "stage": "tool_output", "type": "llm_judge",
1084 "prompt": "Block responses containing PII.", "on_fail": "log"},
1085 ]
1086 }))
1087 .expect("compiles");
1088 assert!(compiled.has_stage(GuardrailStage::ToolUse));
1089 assert!(compiled.has_stage(GuardrailStage::ToolOutput));
1090 assert!(
1092 compiled
1093 .evaluate(
1094 GuardrailStage::ToolUse,
1095 "{}",
1096 Some("delete_user"),
1097 &no_skip()
1098 )
1099 .is_empty()
1100 );
1101 let use_checks: Vec<_> = compiled
1102 .judge_checks_for_stage(GuardrailStage::ToolUse)
1103 .collect();
1104 assert_eq!(use_checks.len(), 1);
1105 assert_eq!(use_checks[0].prompt, "Block requests to delete data.");
1106 assert_eq!(use_checks[0].on_fail, GuardrailOnFail::Block);
1107
1108 let out_checks: Vec<_> = compiled
1109 .judge_checks_for_stage(GuardrailStage::ToolOutput)
1110 .collect();
1111 assert_eq!(out_checks.len(), 1);
1112 assert_eq!(out_checks[0].label, "tj2");
1113 assert_eq!(out_checks[0].on_fail, GuardrailOnFail::Log);
1114 }
1115
1116 #[test]
1117 fn llm_judge_rejected_on_output_stage() {
1118 let err = compile(json!({
1119 "checks": [
1120 {"stage": "output", "type": "llm_judge", "prompt": "Block bad content."},
1121 ]
1122 }))
1123 .unwrap_err();
1124 assert!(err.contains("not supported on the 'output' stage"), "{err}");
1125 }
1126
1127 #[test]
1128 fn llm_judge_empty_prompt_rejected() {
1129 let err = compile(json!({
1130 "checks": [
1131 {"stage": "tool_use", "type": "llm_judge", "prompt": ""},
1132 ]
1133 }))
1134 .unwrap_err();
1135 assert!(err.contains("prompt must not be empty"), "{err}");
1136 }
1137
1138 #[test]
1139 fn llm_judge_prompt_too_long_rejected() {
1140 let long_prompt = "x".repeat(MAX_JUDGE_PROMPT_LEN + 1);
1141 let err = compile(json!({
1142 "checks": [
1143 {"stage": "tool_use", "type": "llm_judge", "prompt": long_prompt},
1144 ]
1145 }))
1146 .unwrap_err();
1147 assert!(err.contains("exceeds"), "{err}");
1148 }
1149
1150 #[test]
1151 fn llm_judge_not_in_sync_evaluate() {
1152 let compiled = compile(json!({
1155 "checks": [
1156 {"stage": "tool_use", "type": "llm_judge", "prompt": "Block everything."},
1157 ]
1158 }))
1159 .unwrap();
1160 assert!(
1161 compiled
1162 .evaluate(
1163 GuardrailStage::ToolUse,
1164 "anything",
1165 Some("tool"),
1166 &no_skip()
1167 )
1168 .is_empty()
1169 );
1170 }
1171
1172 #[test]
1173 fn llm_judge_advisory_downgrades_judge_action() {
1174 let compiled = compile(json!({
1175 "mode": "advisory",
1176 "checks": [
1177 {"stage": "tool_use", "type": "llm_judge", "prompt": "p", "on_fail": "block"},
1178 ]
1179 }))
1180 .unwrap();
1181 let check = compiled
1182 .judge_checks_for_stage(GuardrailStage::ToolUse)
1183 .next()
1184 .unwrap();
1185 assert_eq!(compiled.judge_action(check.on_fail), GuardrailAction::Log);
1187 }
1188
1189 #[test]
1190 fn llm_judge_active_block_yields_block_action() {
1191 let compiled = compile(json!({
1192 "checks": [
1193 {"stage": "tool_use", "type": "llm_judge", "prompt": "p", "on_fail": "block"},
1194 ]
1195 }))
1196 .unwrap();
1197 let check = compiled
1198 .judge_checks_for_stage(GuardrailStage::ToolUse)
1199 .next()
1200 .unwrap();
1201 assert_eq!(compiled.judge_action(check.on_fail), GuardrailAction::Block);
1202 }
1203
1204 #[test]
1205 fn llm_judge_serde_roundtrip() {
1206 let cfg = GuardrailsConfig {
1207 mode: GuardrailMode::Active,
1208 checks: vec![GuardrailCheck {
1209 id: Some("pii-judge".into()),
1210 stage: GuardrailStage::ToolOutput,
1211 on_fail: GuardrailOnFail::Log,
1212 replacement: None,
1213 rule: GuardrailRule::LlmJudge {
1214 prompt: "Block responses that contain PII.".into(),
1215 },
1216 }],
1217 };
1218 let value = serde_json::to_value(&cfg).unwrap();
1219 assert_eq!(value["checks"][0]["type"], "llm_judge");
1220 assert_eq!(value["checks"][0]["stage"], "tool_output");
1221 assert_eq!(
1222 value["checks"][0]["prompt"],
1223 "Block responses that contain PII."
1224 );
1225 let back = GuardrailsConfig::from_value(&value).unwrap();
1226 assert_eq!(back, cfg);
1227 }
1228
1229 #[test]
1230 fn mixed_sync_and_judge_checks_compile_independently() {
1231 let compiled = compile(json!({
1234 "checks": [
1235 {"stage": "tool_use", "type": "tool_pattern", "tools": ["bash*"]},
1236 {"stage": "tool_use", "type": "llm_judge", "prompt": "Block policy violations."},
1237 ]
1238 }))
1239 .unwrap();
1240 let hits = compiled.evaluate(GuardrailStage::ToolUse, "{}", Some("bash_exec"), &no_skip());
1242 assert_eq!(hits.len(), 1);
1243 assert_eq!(hits[0].rule_type, "tool_pattern");
1244 let judges: Vec<_> = compiled
1246 .judge_checks_for_stage(GuardrailStage::ToolUse)
1247 .collect();
1248 assert_eq!(judges.len(), 1);
1249 }
1250
1251 #[test]
1254 fn mcp_compiles_for_tool_stages() {
1255 let compiled = compile(json!({
1256 "checks": [
1257 {"stage": "tool_use", "type": "mcp", "server": "guard", "tool": "screen"},
1258 {"id": "mc2", "stage": "tool_output", "type": "mcp",
1259 "server": "guard", "tool": "scan", "on_fail": "log"},
1260 ]
1261 }))
1262 .expect("compiles");
1263 assert!(compiled.has_stage(GuardrailStage::ToolUse));
1264 assert!(compiled.has_stage(GuardrailStage::ToolOutput));
1265 assert!(
1267 compiled
1268 .evaluate(
1269 GuardrailStage::ToolUse,
1270 "{}",
1271 Some("delete_user"),
1272 &no_skip()
1273 )
1274 .is_empty()
1275 );
1276 let use_checks: Vec<_> = compiled
1277 .mcp_checks_for_stage(GuardrailStage::ToolUse)
1278 .collect();
1279 assert_eq!(use_checks.len(), 1);
1280 assert_eq!(use_checks[0].server, "guard");
1281 assert_eq!(use_checks[0].tool, "screen");
1282 assert_eq!(use_checks[0].on_fail, GuardrailOnFail::Block);
1283
1284 let out_checks: Vec<_> = compiled
1285 .mcp_checks_for_stage(GuardrailStage::ToolOutput)
1286 .collect();
1287 assert_eq!(out_checks.len(), 1);
1288 assert_eq!(out_checks[0].label, "mc2");
1289 assert_eq!(out_checks[0].on_fail, GuardrailOnFail::Log);
1290 }
1291
1292 #[test]
1293 fn mcp_rejected_on_output_stage() {
1294 let err = compile(json!({
1295 "checks": [
1296 {"stage": "output", "type": "mcp", "server": "guard", "tool": "scan"},
1297 ]
1298 }))
1299 .unwrap_err();
1300 assert!(err.contains("not supported on the 'output' stage"), "{err}");
1301 }
1302
1303 #[test]
1304 fn mcp_empty_server_or_tool_rejected() {
1305 let err = compile(json!({
1306 "checks": [
1307 {"stage": "tool_use", "type": "mcp", "server": "", "tool": "scan"},
1308 ]
1309 }))
1310 .unwrap_err();
1311 assert!(err.contains("server must not be empty"), "{err}");
1312 let err = compile(json!({
1313 "checks": [
1314 {"stage": "tool_use", "type": "mcp", "server": "guard", "tool": ""},
1315 ]
1316 }))
1317 .unwrap_err();
1318 assert!(err.contains("tool must not be empty"), "{err}");
1319 }
1320
1321 #[test]
1322 fn mcp_ref_too_long_rejected() {
1323 let long = "x".repeat(MAX_MCP_REF_LEN + 1);
1324 let err = compile(json!({
1325 "checks": [
1326 {"stage": "tool_use", "type": "mcp", "server": long, "tool": "scan"},
1327 ]
1328 }))
1329 .unwrap_err();
1330 assert!(err.contains("exceeds"), "{err}");
1331 }
1332
1333 #[test]
1334 fn mcp_not_in_sync_evaluate() {
1335 let compiled = compile(json!({
1336 "checks": [
1337 {"stage": "tool_use", "type": "mcp", "server": "guard", "tool": "scan"},
1338 ]
1339 }))
1340 .unwrap();
1341 assert!(
1342 compiled
1343 .evaluate(
1344 GuardrailStage::ToolUse,
1345 "anything",
1346 Some("tool"),
1347 &no_skip()
1348 )
1349 .is_empty()
1350 );
1351 }
1352
1353 #[test]
1354 fn mcp_advisory_downgrades_action() {
1355 let compiled = compile(json!({
1356 "mode": "advisory",
1357 "checks": [
1358 {"stage": "tool_use", "type": "mcp", "server": "g", "tool": "t", "on_fail": "block"},
1359 ]
1360 }))
1361 .unwrap();
1362 let check = compiled
1363 .mcp_checks_for_stage(GuardrailStage::ToolUse)
1364 .next()
1365 .unwrap();
1366 assert_eq!(compiled.async_action(check.on_fail), GuardrailAction::Log);
1367 }
1368
1369 #[test]
1370 fn mcp_active_block_yields_block_action() {
1371 let compiled = compile(json!({
1372 "checks": [
1373 {"stage": "tool_use", "type": "mcp", "server": "g", "tool": "t", "on_fail": "block"},
1374 ]
1375 }))
1376 .unwrap();
1377 let check = compiled
1378 .mcp_checks_for_stage(GuardrailStage::ToolUse)
1379 .next()
1380 .unwrap();
1381 assert_eq!(compiled.async_action(check.on_fail), GuardrailAction::Block);
1382 }
1383
1384 #[test]
1385 fn mcp_serde_roundtrip() {
1386 let cfg = GuardrailsConfig {
1387 mode: GuardrailMode::Active,
1388 checks: vec![GuardrailCheck {
1389 id: Some("ext-guard".into()),
1390 stage: GuardrailStage::ToolOutput,
1391 on_fail: GuardrailOnFail::Log,
1392 replacement: None,
1393 rule: GuardrailRule::Mcp {
1394 server: "guard".into(),
1395 tool: "scan".into(),
1396 },
1397 }],
1398 };
1399 let value = serde_json::to_value(&cfg).unwrap();
1400 assert_eq!(value["checks"][0]["type"], "mcp");
1401 assert_eq!(value["checks"][0]["stage"], "tool_output");
1402 assert_eq!(value["checks"][0]["server"], "guard");
1403 assert_eq!(value["checks"][0]["tool"], "scan");
1404 let back = GuardrailsConfig::from_value(&value).unwrap();
1405 assert_eq!(back, cfg);
1406 }
1407}