agentic_warden/
task_record.rs

1use crate::core::models::{AiCliProcessInfo, ProcessTreeInfo};
2use crate::error::AgenticResult;
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, schemars::JsonSchema)]
7#[serde(rename_all = "snake_case")]
8pub enum TaskStatus {
9    #[default]
10    Running,
11    CompletedButUnread,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct TaskRecord {
16    pub started_at: DateTime<Utc>,
17    pub log_id: String,
18    pub log_path: String,
19    #[serde(default)]
20    pub manager_pid: Option<u32>,
21    #[serde(default)]
22    pub cleanup_reason: Option<String>,
23    #[serde(default)]
24    pub status: TaskStatus,
25    #[serde(default)]
26    pub result: Option<String>,
27    #[serde(default)]
28    pub completed_at: Option<DateTime<Utc>>,
29    #[serde(default)]
30    pub exit_code: Option<i32>,
31    // New fields for process tree tracking
32    #[serde(default)]
33    pub process_chain: Vec<u32>,
34    #[serde(default)]
35    pub root_parent_pid: Option<u32>,
36    #[serde(default)]
37    pub process_tree_depth: usize,
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub process_tree: Option<ProcessTreeInfo>,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub ai_cli_process: Option<AiCliProcessInfo>,
42}
43
44impl TaskRecord {
45    pub fn new(
46        started_at: DateTime<Utc>,
47        log_id: String,
48        log_path: String,
49        manager_pid: Option<u32>,
50    ) -> Self {
51        Self {
52            started_at,
53            log_id,
54            log_path,
55            manager_pid,
56            cleanup_reason: None,
57            status: TaskStatus::Running,
58            result: None,
59            completed_at: None,
60            exit_code: None,
61            process_chain: Vec::new(),
62            root_parent_pid: None,
63            process_tree_depth: 0,
64            process_tree: None,
65            ai_cli_process: None,
66        }
67    }
68
69    pub fn with_process_tree_info(mut self, process_tree: ProcessTreeInfo) -> AgenticResult<Self> {
70        process_tree.validate()?;
71        self.process_chain = process_tree.process_chain.clone();
72        self.root_parent_pid = process_tree.root_parent_pid;
73        self.process_tree_depth = process_tree.depth;
74        self.ai_cli_process = process_tree.ai_cli_process.clone();
75        self.process_tree = Some(process_tree);
76        Ok(self)
77    }
78
79    pub fn resolved_root_parent_pid(&self) -> Option<u32> {
80        self.process_tree
81            .as_ref()
82            .and_then(|tree| tree.get_ai_cli_root())
83            .or(self.root_parent_pid)
84    }
85
86    pub fn mark_completed(
87        mut self,
88        result: Option<String>,
89        exit_code: Option<i32>,
90        completed_at: DateTime<Utc>,
91    ) -> Self {
92        self.status = TaskStatus::CompletedButUnread;
93        self.result = result;
94        self.exit_code = exit_code;
95        self.completed_at = Some(completed_at);
96        self
97    }
98
99    pub fn with_cleanup_reason(mut self, reason: &str) -> Self {
100        let result = self.result.clone();
101        let exit_code = self.exit_code;
102        let completed_at = self.completed_at.unwrap_or_else(Utc::now);
103        self = self.mark_completed(result, exit_code, completed_at);
104        self.cleanup_reason = Some(reason.to_owned());
105        self
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use chrono::Utc;
113
114    #[test]
115    fn test_task_record_with_process_tree() {
116        let base_time = Utc::now();
117        let record = TaskRecord::new(
118            base_time,
119            "1234".to_string(),
120            "/tmp/1234.log".to_string(),
121            Some(5678),
122        );
123
124        let process_chain = vec![5678, 1, 0];
125        let tree_info = ProcessTreeInfo::new(process_chain.clone());
126        let depth = tree_info.depth;
127        let root_parent_pid = tree_info.root_parent_pid;
128
129        let enhanced_record = record
130            .with_process_tree_info(tree_info.clone())
131            .expect("process tree injection should succeed");
132
133        assert_eq!(enhanced_record.process_chain, process_chain);
134        assert_eq!(enhanced_record.root_parent_pid, root_parent_pid);
135        assert_eq!(enhanced_record.process_tree_depth, depth);
136        assert!(enhanced_record.process_tree.is_some());
137        assert_eq!(enhanced_record.resolved_root_parent_pid(), root_parent_pid);
138        assert_eq!(enhanced_record.log_id, "1234");
139        assert_eq!(enhanced_record.manager_pid, Some(5678));
140    }
141
142    #[test]
143    fn test_task_record_serialization_with_process_tree() {
144        let base_time = Utc::now();
145        let record = TaskRecord::new(
146            base_time,
147            "1234".to_string(),
148            "/tmp/1234.log".to_string(),
149            Some(5678),
150        );
151        let record = record
152            .with_process_tree_info(ProcessTreeInfo::new(vec![5678, 1]))
153            .expect("process tree should attach");
154
155        // Test that the record can be serialized to JSON
156        let json_str = serde_json::to_string(&record).expect("Failed to serialize");
157
158        // Test that it can be deserialized back
159        let deserialized: TaskRecord =
160            serde_json::from_str(&json_str).expect("Failed to deserialize");
161
162        assert_eq!(deserialized.process_chain, vec![5678, 1]);
163        assert_eq!(deserialized.root_parent_pid, Some(1));
164        assert_eq!(deserialized.process_tree_depth, 2);
165        assert!(deserialized.process_tree.is_some());
166    }
167
168    #[test]
169    fn test_task_record_backward_compatibility() {
170        // Test that a record without process tree fields can still be deserialized
171        let old_record_json = r#"{
172            "started_at": "2024-01-01T12:00:00Z",
173            "log_id": "1234",
174            "log_path": "/tmp/1234.log",
175            "manager_pid": 5678,
176            "cleanup_reason": null,
177            "status": "running",
178            "result": null,
179            "completed_at": null,
180            "exit_code": null
181        }"#;
182
183        let deserialized: TaskRecord =
184            serde_json::from_str(old_record_json).expect("Failed to deserialize old format");
185
186        assert_eq!(deserialized.process_chain, Vec::<u32>::new());
187        assert_eq!(deserialized.root_parent_pid, None);
188        assert_eq!(deserialized.process_tree_depth, 0);
189        assert_eq!(deserialized.log_id, "1234");
190        assert_eq!(deserialized.manager_pid, Some(5678));
191    }
192
193    #[test]
194    fn test_task_record_mark_completed_preserves_process_tree() {
195        let base_time = Utc::now();
196        let record = TaskRecord::new(
197            base_time,
198            "1234".to_string(),
199            "/tmp/1234.log".to_string(),
200            Some(5678),
201        )
202        .with_process_tree_info(ProcessTreeInfo::new(vec![5678, 1]))
203        .expect("process tree should attach");
204
205        let completed_record =
206            record.mark_completed(Some("success".to_string()), Some(0), Utc::now());
207
208        // Process tree fields should be preserved
209        assert_eq!(completed_record.process_chain, vec![5678, 1]);
210        assert_eq!(completed_record.root_parent_pid, Some(1));
211        assert_eq!(completed_record.process_tree_depth, 2);
212        assert!(completed_record.process_tree.is_some());
213
214        // Status should be updated
215        assert_eq!(completed_record.status, TaskStatus::CompletedButUnread);
216        assert_eq!(completed_record.result, Some("success".to_string()));
217        assert_eq!(completed_record.exit_code, Some(0));
218    }
219}
agentic_warden/task_record.rs

agentic_warden/
task_record.rs
