bamboo_agent/agent/loop_module/
todo_context.rs

1//! TodoList context for Agent Loop integration
2//!
3//! This module provides TodoLoopContext which integrates TodoList
4//! as a first-class citizen in the Agent Loop, similar to Token Budget.
5
6use crate::agent::core::todo::{TodoItem, TodoItemStatus, TodoList};
7use crate::agent::core::tools::ToolResult;
8use chrono::{DateTime, Utc};
9use serde::{Deserialize, Serialize};
10
11/// TodoList context for Agent Loop
12///
13/// Acts as a first-class citizen in the agent loop, tracking
14/// task progress throughout the entire conversation lifecycle.
15#[derive(Debug, Clone)]
16pub struct TodoLoopContext {
17    /// Session ID
18    pub session_id: String,
19
20    /// Todo items with execution tracking
21    pub items: Vec<TodoLoopItem>,
22
23    /// Currently active todo item ID
24    pub active_item_id: Option<String>,
25
26    /// Current round number
27    pub current_round: u32,
28
29    /// Maximum rounds allowed
30    pub max_rounds: u32,
31
32    /// Creation timestamp
33    pub created_at: DateTime<Utc>,
34
35    /// Last update timestamp
36    pub updated_at: DateTime<Utc>,
37
38    /// Version number for conflict detection
39    pub version: u64,
40}
41
42/// Todo item with execution tracking
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct TodoLoopItem {
45    /// Item ID
46    pub id: String,
47
48    /// Item description
49    pub description: String,
50
51    /// Item status
52    pub status: TodoItemStatus,
53
54    /// Tool call history (tracks execution process)
55    pub tool_calls: Vec<ToolCallRecord>,
56
57    /// Round when item was started
58    pub started_at_round: Option<u32>,
59
60    /// Round when item was completed
61    pub completed_at_round: Option<u32>,
62}
63
64/// Record of a tool call execution
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct ToolCallRecord {
67    /// Round number
68    pub round: u32,
69
70    /// Tool name
71    pub tool_name: String,
72
73    /// Whether the call succeeded
74    pub success: bool,
75
76    /// Timestamp
77    pub timestamp: DateTime<Utc>,
78}
79
80impl TodoLoopContext {
81    /// Create TodoLoopContext from Session's TodoList
82    pub fn from_session(session: &crate::agent::core::Session) -> Option<Self> {
83        session.todo_list.as_ref().map(|todo_list| {
84            // Preserve version from existing todo_list metadata if available
85            // This prevents version reset across multiple executions
86            let existing_version = session
87                .metadata
88                .get("todo_list_version")
89                .and_then(|v| v.parse::<u64>().ok())
90                .unwrap_or(0);
91
92            Self {
93                session_id: todo_list.session_id.clone(),
94                items: todo_list
95                    .items
96                    .iter()
97                    .map(|item| TodoLoopItem {
98                        id: item.id.clone(),
99                        description: item.description.clone(),
100                        status: item.status.clone(),
101                        tool_calls: Vec::new(),
102                        started_at_round: None,
103                        completed_at_round: None,
104                    })
105                    .collect(),
106                active_item_id: None,
107                current_round: 0,
108                max_rounds: 50,
109                created_at: todo_list.created_at,
110                updated_at: todo_list.updated_at,
111                version: existing_version,
112            }
113        })
114    }
115
116    /// Track tool execution
117    ///
118    /// Records a tool call and associates it with the active todo item.
119    pub fn track_tool_execution(&mut self, tool_name: &str, result: &ToolResult, round: u32) {
120        self.current_round = round;
121
122        // Record tool call
123        let record = ToolCallRecord {
124            round,
125            tool_name: tool_name.to_string(),
126            success: result.success,
127            timestamp: Utc::now(),
128        };
129
130        // Associate with active item if exists
131        if let Some(ref active_id) = self.active_item_id {
132            if let Some(item) = self.items.iter_mut().find(|i| &i.id == active_id) {
133                item.tool_calls.push(record);
134                self.updated_at = Utc::now();
135                self.version += 1;
136            }
137        }
138    }
139
140    /// Set active todo item
141    ///
142    /// Marks the previous active item as completed and activates a new item.
143    pub fn set_active_item(&mut self, item_id: &str) {
144        // Complete previous active item
145        if let Some(ref prev_id) = self.active_item_id {
146            if let Some(item) = self.items.iter_mut().find(|i| &i.id == prev_id) {
147                item.status = TodoItemStatus::Completed;
148                item.completed_at_round = Some(self.current_round);
149            }
150        }
151
152        // Set new active item
153        self.active_item_id = Some(item_id.to_string());
154        if let Some(item) = self.items.iter_mut().find(|i| i.id == item_id) {
155            item.status = TodoItemStatus::InProgress;
156            item.started_at_round = Some(self.current_round);
157        }
158
159        self.updated_at = Utc::now();
160        self.version += 1;
161    }
162
163    /// Update item status manually
164    pub fn update_item_status(&mut self, item_id: &str, status: TodoItemStatus) {
165        if let Some(item) = self.items.iter_mut().find(|i| i.id == item_id) {
166            item.status = status.clone();
167
168            match &status {
169                TodoItemStatus::InProgress => {
170                    item.started_at_round = Some(self.current_round);
171                    self.active_item_id = Some(item_id.to_string());
172                }
173                TodoItemStatus::Completed => {
174                    item.completed_at_round = Some(self.current_round);
175                    if self.active_item_id.as_deref() == Some(item_id) {
176                        self.active_item_id = None;
177                    }
178                }
179                _ => {}
180            }
181
182            self.updated_at = Utc::now();
183            self.version += 1;
184        }
185    }
186
187    /// Check if all items are completed
188    pub fn is_all_completed(&self) -> bool {
189        !self.items.is_empty()
190            && self
191                .items
192                .iter()
193                .all(|item| matches!(item.status, TodoItemStatus::Completed))
194    }
195
196    /// Generate context for prompt injection
197    pub fn format_for_prompt(&self) -> String {
198        if self.items.is_empty() {
199            return String::new();
200        }
201
202        let mut output = format!(
203            "\n\n## Current Task List (Round {}/{})\n",
204            self.current_round + 1,
205            self.max_rounds
206        );
207
208        for item in &self.items {
209            let status_icon = match item.status {
210                TodoItemStatus::Pending => "[ ]",
211                TodoItemStatus::InProgress => "[/]",
212                TodoItemStatus::Completed => "[x]",
213                TodoItemStatus::Blocked => "[!]",
214            };
215
216            output.push_str(&format!(
217                "\n{} {}: {}",
218                status_icon, item.id, item.description
219            ));
220
221            if !item.tool_calls.is_empty() {
222                output.push_str(&format!(" ({} tool calls)", item.tool_calls.len()));
223            }
224        }
225
226        let completed = self
227            .items
228            .iter()
229            .filter(|i| matches!(i.status, TodoItemStatus::Completed))
230            .count();
231        output.push_str(&format!(
232            "\n\nProgress: {}/{} tasks completed",
233            completed,
234            self.items.len()
235        ));
236
237        output
238    }
239
240    /// Convert back to TodoList for persistence
241    pub fn into_todo_list(self) -> TodoList {
242        TodoList {
243            session_id: self.session_id,
244            title: "Agent Tasks".to_string(),
245            items: self
246                .items
247                .into_iter()
248                .map(|loop_item| TodoItem {
249                    id: loop_item.id,
250                    description: loop_item.description,
251                    status: loop_item.status,
252                    depends_on: Vec::new(),
253                    notes: String::new(),
254                })
255                .collect(),
256            created_at: self.created_at,
257            updated_at: self.updated_at,
258        }
259    }
260
261    /// Auto-match tool to todo item based on keywords
262    pub fn auto_match_tool_to_item(&mut self, tool_name: &str) {
263        if self.active_item_id.is_some() {
264            return; // Already have an active item
265        }
266
267        // Strategy: Keyword matching
268        let tool_lower = tool_name.to_lowercase();
269        let matching_item_id = self.items.iter().find(|item| {
270            let desc_lower = item.description.to_lowercase();
271            // Simple heuristic: tool name appears in description
272            desc_lower.contains(&tool_lower) ||
273            // Or common tool patterns
274            (tool_lower.contains("file") && desc_lower.contains("file")) ||
275            (tool_lower.contains("command") && (desc_lower.contains("run") || desc_lower.contains("execute")))
276        }).map(|item| item.id.clone());
277
278        if let Some(item_id) = matching_item_id {
279            self.set_active_item(&item_id);
280        }
281    }
282
283    /// Auto-update status based on tool execution result
284    pub fn auto_update_status(&mut self, tool_name: &str, result: &ToolResult) {
285        // If no active item, try to auto-match
286        if self.active_item_id.is_none() {
287            self.auto_match_tool_to_item(tool_name);
288        }
289
290        if let Some(ref active_id) = self.active_item_id.clone() {
291            // First, determine what action to take (avoid borrow issues)
292            let action = self
293                .items
294                .iter()
295                .find(|i| &i.id == active_id)
296                .and_then(|item| {
297                    if result.success {
298                        if self.should_mark_completed(item) {
299                            Some(TodoItemStatus::Completed)
300                        } else {
301                            None
302                        }
303                    } else if self.should_mark_blocked(item) {
304                        Some(TodoItemStatus::Blocked)
305                    } else {
306                        None
307                    }
308                });
309
310            // Then apply the action
311            if let Some(new_status) = action {
312                if let Some(item) = self.items.iter_mut().find(|i| &i.id == active_id) {
313                    item.status = new_status.clone();
314                    if matches!(new_status, TodoItemStatus::Completed) {
315                        item.completed_at_round = Some(self.current_round);
316                        self.active_item_id = None;
317                    }
318                    self.version += 1;
319                    self.updated_at = Utc::now(); // IMPORTANT: Update timestamp
320                }
321            }
322        }
323    }
324
325    /// Determine if item should be marked as completed
326    fn should_mark_completed(&self, item: &TodoLoopItem) -> bool {
327        // Simple strategy: 3 successful tool calls
328        let success_count = item.tool_calls.iter().filter(|r| r.success).count();
329        success_count >= 3
330    }
331
332    /// Determine if item should be marked as blocked
333    fn should_mark_blocked(&self, item: &TodoLoopItem) -> bool {
334        // Simple strategy: 2 consecutive failures
335        let recent_failures = item
336            .tool_calls
337            .iter()
338            .rev()
339            .take(2)
340            .filter(|r| !r.success)
341            .count();
342        recent_failures >= 2
343    }
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349
350    fn create_test_todo_list() -> crate::agent::core::Session {
351        let mut session = crate::agent::core::Session::new("test-session", "test-model");
352        let todo_list = TodoList {
353            session_id: "test-session".to_string(),
354            title: "Test Tasks".to_string(),
355            items: vec![
356                TodoItem {
357                    id: "task-1".to_string(),
358                    description: "Read configuration file".to_string(),
359                    status: TodoItemStatus::Pending,
360                    depends_on: Vec::new(),
361                    notes: String::new(),
362                },
363                TodoItem {
364                    id: "task-2".to_string(),
365                    description: "Run tests".to_string(),
366                    status: TodoItemStatus::Pending,
367                    depends_on: Vec::new(),
368                    notes: String::new(),
369                },
370            ],
371            created_at: Utc::now(),
372            updated_at: Utc::now(),
373        };
374        session.set_todo_list(todo_list);
375        session
376    }
377
378    #[test]
379    fn test_from_session() {
380        let session = create_test_todo_list();
381        let ctx = TodoLoopContext::from_session(&session).unwrap();
382
383        assert_eq!(ctx.session_id, "test-session");
384        assert_eq!(ctx.items.len(), 2);
385        assert_eq!(ctx.items[0].id, "task-1");
386        assert_eq!(ctx.items[0].tool_calls.len(), 0);
387    }
388
389    #[test]
390    fn test_track_tool_execution() {
391        let session = create_test_todo_list();
392        let mut ctx = TodoLoopContext::from_session(&session).unwrap();
393
394        // Set active item
395        ctx.set_active_item("task-1");
396
397        // Track tool execution
398        let result = ToolResult {
399            success: true,
400            result: "OK".to_string(),
401            display_preference: None,
402        };
403        ctx.track_tool_execution("read_file", &result, 1);
404
405        assert_eq!(ctx.items[0].tool_calls.len(), 1);
406        assert_eq!(ctx.items[0].tool_calls[0].tool_name, "read_file");
407        assert!(ctx.items[0].tool_calls[0].success);
408        assert_eq!(ctx.version, 2); // set_active_item + track_tool_execution
409    }
410
411    #[test]
412    fn test_set_active_item() {
413        let session = create_test_todo_list();
414        let mut ctx = TodoLoopContext::from_session(&session).unwrap();
415
416        ctx.set_active_item("task-1");
417
418        assert_eq!(ctx.active_item_id, Some("task-1".to_string()));
419        assert_eq!(ctx.items[0].status, TodoItemStatus::InProgress);
420        assert_eq!(ctx.items[0].started_at_round, Some(0));
421    }
422
423    #[test]
424    fn test_is_all_completed() {
425        let session = create_test_todo_list();
426        let mut ctx = TodoLoopContext::from_session(&session).unwrap();
427
428        assert!(!ctx.is_all_completed());
429
430        ctx.items[0].status = TodoItemStatus::Completed;
431        ctx.items[1].status = TodoItemStatus::Completed;
432
433        assert!(ctx.is_all_completed());
434    }
435
436    #[test]
437    fn test_format_for_prompt() {
438        let session = create_test_todo_list();
439        let mut ctx = TodoLoopContext::from_session(&session).unwrap();
440        ctx.current_round = 2;
441        ctx.max_rounds = 10;
442
443        let prompt = ctx.format_for_prompt();
444
445        assert!(prompt.contains("Round 3/10"));
446        assert!(prompt.contains("task-1"));
447        assert!(prompt.contains("task-2"));
448    }
449
450    #[test]
451    fn test_auto_match_tool_to_item() {
452        let session = create_test_todo_list();
453        let mut ctx = TodoLoopContext::from_session(&session).unwrap();
454
455        // Test keyword matching
456        ctx.auto_match_tool_to_item("read_file");
457
458        // Should match task-1 (Read configuration file)
459        assert_eq!(ctx.active_item_id, Some("task-1".to_string()));
460    }
461
462    #[test]
463    fn test_auto_update_status_completed() {
464        let session = create_test_todo_list();
465        let mut ctx = TodoLoopContext::from_session(&session).unwrap();
466
467        ctx.set_active_item("task-1");
468        ctx.current_round = 1;
469
470        // Add 3 successful tool calls
471        let result = ToolResult {
472            success: true,
473            result: "OK".to_string(),
474            display_preference: None,
475        };
476
477        ctx.track_tool_execution("read_file", &result, 1);
478        ctx.auto_update_status("read_file", &result);
479        assert_eq!(ctx.items[0].status, TodoItemStatus::InProgress);
480
481        ctx.track_tool_execution("read_file", &result, 2);
482        ctx.auto_update_status("read_file", &result);
483        assert_eq!(ctx.items[0].status, TodoItemStatus::InProgress);
484
485        ctx.track_tool_execution("read_file", &result, 3);
486        ctx.auto_update_status("read_file", &result);
487
488        // After 3 successful calls, should be completed
489        assert_eq!(ctx.items[0].status, TodoItemStatus::Completed);
490        assert_eq!(ctx.active_item_id, None);
491    }
492
493    #[test]
494    fn test_auto_update_status_blocked() {
495        let session = create_test_todo_list();
496        let mut ctx = TodoLoopContext::from_session(&session).unwrap();
497
498        ctx.set_active_item("task-1");
499        ctx.current_round = 1;
500
501        // Add 2 failed tool calls
502        let fail_result = ToolResult {
503            success: false,
504            result: "Error".to_string(),
505            display_preference: None,
506        };
507
508        ctx.track_tool_execution("read_file", &fail_result, 1);
509        ctx.auto_update_status("read_file", &fail_result);
510        assert_eq!(ctx.items[0].status, TodoItemStatus::InProgress);
511
512        ctx.track_tool_execution("read_file", &fail_result, 2);
513        ctx.auto_update_status("read_file", &fail_result);
514
515        // After 2 consecutive failures, should be blocked
516        assert_eq!(ctx.items[0].status, TodoItemStatus::Blocked);
517    }
518
519    #[test]
520    fn test_into_todo_list() {
521        let session = create_test_todo_list();
522        let ctx = TodoLoopContext::from_session(&session).unwrap();
523
524        let todo_list = ctx.into_todo_list();
525
526        assert_eq!(todo_list.session_id, "test-session");
527        assert_eq!(todo_list.items.len(), 2);
528    }
529}
bamboo_agent/agent/loop_module/todo_context.rs

bamboo_agent/agent/loop_module/
todo_context.rs
