1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4
5use crate::types::AgentResult;
6
7#[cfg(test)]
8use serde_json::json;
9
10#[async_trait]
11pub trait ReflexionHandler: Send + Sync {
12 async fn reflect_on_failure(
13 &self,
14 failed_action: &str,
15 error: &str,
16 context: &str,
17 ) -> AgentResult<ReflectionResult>;
18
19 async fn generate_alternatives(
20 &self,
21 reflection: &ReflectionResult,
22 ) -> AgentResult<Vec<AlternativeAction>>;
23
24 async fn should_retry(
25 &self,
26 reflection: &ReflectionResult,
27 retry_count: usize,
28 ) -> AgentResult<bool>;
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct ReflectionResult {
33 pub analysis: String,
34 pub root_cause: String,
35 pub confidence: f32,
36 pub suggested_fixes: Vec<String>,
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct AlternativeAction {
41 pub description: String,
42 pub payload: Value,
43 pub priority: u32,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ReflexionContext {
48 pub objective: String,
49 pub failed_step_id: String,
50 pub failed_step_description: String,
51 pub step_payload: Value,
52 pub error: String,
53 pub previous_steps: Vec<StepHistoryEntry>,
54}
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct StepHistoryEntry {
58 pub step_id: String,
59 pub description: String,
60 pub success: bool,
61 pub output: Option<String>,
62}
63
64impl ReflectionResult {
65 pub fn new(
66 analysis: impl Into<String>,
67 root_cause: impl Into<String>,
68 confidence: f32,
69 suggested_fixes: Vec<String>,
70 ) -> Self {
71 Self {
72 analysis: analysis.into(),
73 root_cause: root_cause.into(),
74 confidence,
75 suggested_fixes,
76 }
77 }
78
79 pub fn is_confident(&self) -> bool {
80 self.confidence >= 0.7
81 }
82}
83
84impl AlternativeAction {
85 pub fn new(
86 description: impl Into<String>,
87 payload: Value,
88 priority: u32,
89 ) -> Self {
90 Self {
91 description: description.into(),
92 payload,
93 priority,
94 }
95 }
96}
97
98impl ReflexionContext {
99 pub fn from_step(
100 objective: &str,
101 step: &crate::types::PlanStep,
102 error: &str,
103 ) -> Self {
104 Self {
105 objective: objective.to_string(),
106 failed_step_id: step.id.clone(),
107 failed_step_description: step.description.clone(),
108 step_payload: step.payload.clone(),
109 error: error.to_string(),
110 previous_steps: Vec::new(),
111 }
112 }
113
114 pub fn with_previous_steps(mut self, steps: Vec<StepHistoryEntry>) -> Self {
115 self.previous_steps = steps;
116 self
117 }
118}
119
120#[cfg(test)]
121mod tests {
122 use super::*;
123 use crate::types::PlanStep;
124
125 #[test]
126 fn test_reflection_result_new() {
127 let result = ReflectionResult::new("analysis", "root_cause", 0.8, vec!["fix1".to_string()]);
128
129 assert_eq!(result.analysis, "analysis");
130 assert_eq!(result.root_cause, "root_cause");
131 assert_eq!(result.confidence, 0.8);
132 assert_eq!(result.suggested_fixes.len(), 1);
133 }
134
135 #[test]
136 fn test_reflection_result_is_confident() {
137 let high_confidence = ReflectionResult::new("", "", 0.9, vec![]);
138 assert!(high_confidence.is_confident());
139
140 let low_confidence = ReflectionResult::new("", "", 0.5, vec![]);
141 assert!(!low_confidence.is_confident());
142
143 let threshold = ReflectionResult::new("", "", 0.7, vec![]);
144 assert!(threshold.is_confident());
145 }
146
147 #[test]
148 fn test_alternative_action_new() {
149 let action = AlternativeAction::new(
150 "description",
151 json!({"type":"ssh_command","command":"ls"}),
152 1,
153 );
154
155 assert_eq!(action.description, "description");
156 assert_eq!(action.priority, 1);
157 }
158
159 #[test]
160 fn test_reflexion_context_from_step() {
161 let step = PlanStep::new(
162 "step-1",
163 "check disk",
164 json!({"type":"ssh_command","command":"df -h","host_id":"host1"}),
165 );
166
167 let context = ReflexionContext::from_step("objective", &step, "error");
168
169 assert_eq!(context.objective, "objective");
170 assert_eq!(context.failed_step_id, "step-1");
171 assert_eq!(context.failed_step_description, "check disk");
172 assert_eq!(context.step_payload, json!({"type":"ssh_command","command":"df -h","host_id":"host1"}));
173 assert_eq!(context.error, "error");
174 assert!(context.previous_steps.is_empty());
175 }
176
177 #[test]
178 fn test_reflexion_context_with_previous_steps() {
179 let context = ReflexionContext {
180 objective: "objective".to_string(),
181 failed_step_id: "step-2".to_string(),
182 failed_step_description: "description".to_string(),
183 step_payload: json!({"type":"test"}),
184 error: "error".to_string(),
185 previous_steps: Vec::new(),
186 }
187 .with_previous_steps(vec![StepHistoryEntry {
188 step_id: "step-1".to_string(),
189 description: "previous step".to_string(),
190 success: true,
191 output: Some("done".to_string()),
192 }]);
193
194 assert_eq!(context.previous_steps.len(), 1);
195 assert_eq!(context.previous_steps[0].step_id, "step-1");
196 }
197
198 #[test]
199 fn test_reflection_result_serialization() {
200 let result = ReflectionResult::new("analysis", "root_cause", 0.8, vec!["fix1".to_string()]);
201
202 let json = serde_json::to_string(&result).unwrap();
203 let deserialized: ReflectionResult = serde_json::from_str(&json).unwrap();
204
205 assert_eq!(deserialized.analysis, "analysis");
206 assert_eq!(deserialized.root_cause, "root_cause");
207 assert_eq!(deserialized.confidence, 0.8);
208 }
209
210 #[test]
211 fn test_step_history_entry_serialization() {
212 let entry = StepHistoryEntry {
213 step_id: "step-1".to_string(),
214 description: "description".to_string(),
215 success: true,
216 output: Some("output".to_string()),
217 };
218
219 let json = serde_json::to_string(&entry).unwrap();
220 let deserialized: StepHistoryEntry = serde_json::from_str(&json).unwrap();
221
222 assert_eq!(deserialized.step_id, "step-1");
223 assert!(deserialized.success);
224 assert_eq!(deserialized.output, Some("output".to_string()));
225 }
226}