ricecoder_workflows/
safety_constraints.rs1use crate::models::{SafetyViolation, WorkflowStep};
4use std::time::Duration;
5
6#[derive(Debug, Clone)]
8pub struct SafetyConstraints {
9 pub max_timeout_ms: u64,
11 pub max_memory_mb: u64,
13 pub max_cpu_percent: u8,
15 pub max_file_handles: u32,
17}
18
19impl SafetyConstraints {
20 pub fn new() -> Self {
22 Self {
23 max_timeout_ms: 300_000, max_memory_mb: 1024, max_cpu_percent: 80,
26 max_file_handles: 1024,
27 }
28 }
29
30 pub fn with_timeout(timeout_ms: u64) -> Self {
32 Self {
33 max_timeout_ms: timeout_ms,
34 ..Self::new()
35 }
36 }
37
38 pub fn apply_to_step(&self, step: &WorkflowStep, _risk_score: u8) -> Vec<SafetyViolation> {
40 let mut violations = Vec::new();
41
42 if let crate::models::StepType::Command(cmd_step) = &step.step_type {
44 if cmd_step.timeout > self.max_timeout_ms {
45 violations.push(SafetyViolation {
46 step_id: step.id.clone(),
47 violation_type: "timeout_exceeded".to_string(),
48 description: format!(
49 "Step timeout {} ms exceeds maximum {} ms",
50 cmd_step.timeout, self.max_timeout_ms
51 ),
52 });
53 }
54 }
55
56 violations
57 }
58
59 pub fn enforce_timeout(&self, step: &WorkflowStep) -> Duration {
61 match &step.step_type {
62 crate::models::StepType::Command(cmd_step) => {
63 let timeout_ms = cmd_step.timeout.min(self.max_timeout_ms);
64 Duration::from_millis(timeout_ms)
65 }
66 _ => Duration::from_millis(self.max_timeout_ms),
67 }
68 }
69
70 pub fn has_rollback_capability(&self, step: &WorkflowStep) -> bool {
72 matches!(step.on_error, crate::models::ErrorAction::Rollback)
74 }
75}
76
77impl Default for SafetyConstraints {
78 fn default() -> Self {
79 Self::new()
80 }
81}
82
83#[cfg(test)]
84mod tests {
85 use super::*;
86 use crate::models::{CommandStep, ErrorAction, RiskFactors, StepConfig, StepType};
87
88 fn create_command_step(id: &str, timeout: u64) -> WorkflowStep {
89 WorkflowStep {
90 id: id.to_string(),
91 name: format!("Step {}", id),
92 step_type: StepType::Command(CommandStep {
93 command: "test".to_string(),
94 args: vec![],
95 timeout,
96 }),
97 config: StepConfig {
98 config: serde_json::json!({}),
99 },
100 dependencies: Vec::new(),
101 approval_required: false,
102 on_error: ErrorAction::Rollback,
103 risk_score: None,
104 risk_factors: RiskFactors::default(),
105 }
106 }
107
108 #[test]
109 fn test_safety_constraints_default() {
110 let constraints = SafetyConstraints::new();
111 assert_eq!(constraints.max_timeout_ms, 300_000);
112 assert_eq!(constraints.max_memory_mb, 1024);
113 assert_eq!(constraints.max_cpu_percent, 80);
114 }
115
116 #[test]
117 fn test_enforce_timeout_within_limit() {
118 let constraints = SafetyConstraints::new();
119 let step = create_command_step("1", 100_000);
120 let timeout = constraints.enforce_timeout(&step);
121 assert_eq!(timeout.as_millis(), 100_000);
122 }
123
124 #[test]
125 fn test_enforce_timeout_exceeds_limit() {
126 let constraints = SafetyConstraints::new();
127 let step = create_command_step("1", 500_000);
128 let timeout = constraints.enforce_timeout(&step);
129 assert_eq!(timeout.as_millis(), 300_000);
130 }
131
132 #[test]
133 fn test_timeout_violation_detection() {
134 let constraints = SafetyConstraints::new();
135 let step = create_command_step("1", 500_000);
136 let violations = constraints.apply_to_step(&step, 80);
137 assert!(!violations.is_empty());
138 assert_eq!(violations[0].violation_type, "timeout_exceeded");
139 }
140
141 #[test]
142 fn test_rollback_capability() {
143 let constraints = SafetyConstraints::new();
144 let step = create_command_step("1", 100_000);
145 assert!(constraints.has_rollback_capability(&step));
146 }
147}