Skip to main content

anda_engine/extension/
todo.rs

1//! Todo extension primitives.
2//!
3//! This module provides:
4//! - an in-memory todo store shared through [`BaseCtx`] state,
5//! - tool input/output types for reading and updating the list,
6//! - and the public tool entrypoint ([`TodoTool`]).
7//!
8//! The todo list is session-scoped rather than durable. Repeated calls in the
9//! same context tree, including subagents spawned from that session, see the
10//! same ordered task list.
11
12use anda_core::{BoxError, FunctionDefinition, Resource, Tool, ToolOutput};
13use parking_lot::RwLock;
14use serde::{Deserialize, Serialize};
15use serde_json::json;
16use std::{
17    collections::{HashMap, HashSet},
18    sync::Arc,
19};
20
21use crate::{
22    context::BaseCtx,
23    hook::{DynToolHook, ToolHook},
24};
25
26const TODO_STATUS_PENDING: &str = "pending";
27const TODO_STATUS_IN_PROGRESS: &str = "in_progress";
28const TODO_STATUS_COMPLETED: &str = "completed";
29const TODO_STATUS_CANCELLED: &str = "cancelled";
30const TODO_EMPTY_ID: &str = "?";
31const TODO_EMPTY_CONTENT: &str = "(no description)";
32const TODO_ACTIVE_LIST_PREFIX: &str =
33    "[Your active task list was preserved across context compression]";
34const TODO_MARKER_PENDING: &str = "[ ]";
35const TODO_MARKER_IN_PROGRESS: &str = "[>]";
36const TODO_MARKER_COMPLETED: &str = "[x]";
37const TODO_MARKER_CANCELLED: &str = "[~]";
38
39static VALID_STATUSES: &[&str] = &[
40    TODO_STATUS_PENDING,
41    TODO_STATUS_IN_PROGRESS,
42    TODO_STATUS_COMPLETED,
43    TODO_STATUS_CANCELLED,
44];
45
46/// Shared todo session handle stored on [`BaseCtx`].
47#[derive(Clone, Default)]
48pub struct TodoSession {
49    inner: Arc<RwLock<TodoStore>>,
50}
51
52impl TodoSession {
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    pub fn write(&self, todos: Vec<TodoItemInput>, merge: bool) -> Vec<TodoItem> {
58        self.inner.write().write(todos, merge)
59    }
60
61    pub fn snapshot(&self) -> Vec<TodoItem> {
62        self.inner.read().snapshot()
63    }
64
65    pub fn has_items(&self) -> bool {
66        self.inner.read().has_items()
67    }
68
69    pub fn format_for_injection(&self) -> Option<String> {
70        self.inner.read().format_for_injection()
71    }
72}
73
74/// Returns the current todo session for the context, creating one on demand.
75pub fn todo_session(ctx: &BaseCtx) -> TodoSession {
76    if let Some(session) = ctx.get_state::<TodoSession>() {
77        return session;
78    }
79
80    let session = TodoSession::new();
81    let _ = ctx.set_state(session.clone());
82    session
83}
84
85/// In-memory ordered todo store.
86#[derive(Debug, Clone, Default)]
87pub struct TodoStore {
88    items: Vec<TodoItem>,
89}
90
91impl TodoStore {
92    /// Writes todos using replace or merge semantics, then returns the full list.
93    pub fn write(&mut self, todos: Vec<TodoItemInput>, merge: bool) -> Vec<TodoItem> {
94        let todos = Self::dedupe_by_id(todos);
95
96        if !merge {
97            self.items = todos.into_iter().map(TodoItem::from_input).collect();
98            return self.snapshot();
99        }
100
101        let mut existing: HashMap<String, TodoItem> = self
102            .items
103            .iter()
104            .cloned()
105            .map(|item| (item.id.clone(), item))
106            .collect();
107
108        for todo in todos {
109            let item_id = todo.id.trim().to_string();
110            if item_id.is_empty() {
111                continue;
112            }
113
114            if let Some(item) = existing.get_mut(&item_id) {
115                if let Some(content) = todo.content.as_deref().map(str::trim)
116                    && !content.is_empty()
117                {
118                    item.content = content.to_string();
119                }
120
121                if let Some(status) = todo.status.as_deref() {
122                    item.status = normalize_status(Some(status));
123                }
124            } else {
125                let validated = TodoItem::from_input(todo);
126                existing.insert(validated.id.clone(), validated.clone());
127                self.items.push(validated);
128            }
129        }
130
131        let mut seen = HashSet::new();
132        self.items = self
133            .items
134            .iter()
135            .filter_map(|item| {
136                let current = existing
137                    .get(&item.id)
138                    .cloned()
139                    .unwrap_or_else(|| item.clone());
140                if seen.insert(current.id.clone()) {
141                    Some(current)
142                } else {
143                    None
144                }
145            })
146            .collect();
147
148        self.snapshot()
149    }
150
151    /// Returns a copy of the current ordered todo list.
152    pub fn snapshot(&self) -> Vec<TodoItem> {
153        self.items.clone()
154    }
155
156    /// Returns true when the store contains at least one item.
157    pub fn has_items(&self) -> bool {
158        !self.items.is_empty()
159    }
160
161    /// Renders the active todo items for prompt reinjection after compression.
162    pub fn format_for_injection(&self) -> Option<String> {
163        if self.items.is_empty() {
164            return None;
165        }
166
167        let active_items: Vec<&TodoItem> = self
168            .items
169            .iter()
170            .filter(|item| {
171                matches!(
172                    item.status.as_str(),
173                    TODO_STATUS_PENDING | TODO_STATUS_IN_PROGRESS
174                )
175            })
176            .collect();
177
178        if active_items.is_empty() {
179            return None;
180        }
181
182        let mut lines = Vec::with_capacity(active_items.len() + 1);
183        lines.push(TODO_ACTIVE_LIST_PREFIX.to_string());
184        for item in active_items {
185            lines.push(format!(
186                "- {} {}. {} ({})",
187                status_marker(&item.status),
188                item.id,
189                item.content,
190                item.status
191            ));
192        }
193
194        Some(lines.join("\n"))
195    }
196
197    fn dedupe_by_id(todos: Vec<TodoItemInput>) -> Vec<TodoItemInput> {
198        let mut last_index = HashMap::new();
199        for (index, item) in todos.iter().enumerate() {
200            last_index.insert(todo_dedupe_key(item), index);
201        }
202
203        let mut indexes: Vec<usize> = last_index.into_values().collect();
204        indexes.sort_unstable();
205
206        indexes
207            .into_iter()
208            .map(|index| todos[index].clone())
209            .collect()
210    }
211}
212
213/// Arguments for the todo tool.
214#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
215pub struct TodoArgs {
216    /// Task items to write. Omit to read the current list.
217    #[serde(default, skip_serializing_if = "Option::is_none")]
218    pub todos: Option<Vec<TodoItemInput>>,
219    /// Whether writes should merge into the existing list by id.
220    #[serde(default)]
221    pub merge: bool,
222}
223
224/// Input item accepted by the todo tool.
225#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
226pub struct TodoItemInput {
227    /// Unique todo identifier.
228    #[serde(default)]
229    pub id: String,
230    /// Human-readable todo description.
231    #[serde(default)]
232    pub content: Option<String>,
233    /// Current todo status.
234    #[serde(default)]
235    pub status: Option<String>,
236}
237
238/// Normalized todo item returned by the tool.
239#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
240pub struct TodoItem {
241    /// Unique todo identifier.
242    pub id: String,
243    /// Human-readable todo description.
244    pub content: String,
245    /// Current todo status.
246    pub status: String,
247}
248
249impl TodoItem {
250    fn from_input(input: TodoItemInput) -> Self {
251        let id = input.id.trim();
252        let content = input.content.as_deref().unwrap_or_default().trim();
253
254        Self {
255            id: if id.is_empty() {
256                TODO_EMPTY_ID.to_string()
257            } else {
258                id.to_string()
259            },
260            content: if content.is_empty() {
261                TODO_EMPTY_CONTENT.to_string()
262            } else {
263                content.to_string()
264            },
265            status: normalize_status(input.status.as_deref()),
266        }
267    }
268}
269
270/// Summary counts returned with the todo list.
271#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
272pub struct TodoSummary {
273    pub total: usize,
274    pub pending: usize,
275    pub in_progress: usize,
276    pub completed: usize,
277    pub cancelled: usize,
278}
279
280impl TodoSummary {
281    fn from_items(items: &[TodoItem]) -> Self {
282        let mut summary = Self {
283            total: items.len(),
284            ..Default::default()
285        };
286
287        for item in items {
288            match item.status.as_str() {
289                TODO_STATUS_PENDING => summary.pending += 1,
290                TODO_STATUS_IN_PROGRESS => summary.in_progress += 1,
291                TODO_STATUS_COMPLETED => summary.completed += 1,
292                TODO_STATUS_CANCELLED => summary.cancelled += 1,
293                _ => {}
294            }
295        }
296
297        summary
298    }
299}
300
301/// Output returned by the todo tool.
302#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
303pub struct TodoOutput {
304    pub todos: Vec<TodoItem>,
305    pub summary: TodoSummary,
306}
307
308pub type TodoToolHook = DynToolHook<TodoArgs, TodoOutput>;
309
310/// Tool implementation that exposes the session todo list to the agent.
311#[derive(Clone)]
312pub struct TodoTool {
313    description: String,
314}
315
316impl Default for TodoTool {
317    fn default() -> Self {
318        Self::new()
319    }
320}
321
322impl TodoTool {
323    /// Tool name used for registration and function definition.
324    pub const NAME: &'static str = "todo";
325
326    /// Creates a todo tool with the default behavioral guidance.
327    pub fn new() -> Self {
328        Self {
329            description: concat!(
330                "Manage your task list for the current session. Use for complex tasks ",
331                "with 3+ steps or when the user provides multiple tasks. ",
332                "This list is shared across the current agent and its subagents ",
333                "within the same session/context tree. ",
334                "Call with no parameters to read the current list.\n\n",
335                "Writing:\n",
336                "- Provide 'todos' array to create/update items\n",
337                "- merge=false (default): replace the entire list with a fresh plan\n",
338                "- merge=true: update existing items by id, add any new ones\n\n",
339                "Each item: {id: string, content: string, ",
340                "status: pending|in_progress|completed|cancelled}\n",
341                "List order is priority. Only ONE item in_progress at a time.\n",
342                "Mark items completed immediately when done. If something fails, ",
343                "cancel it and add a revised item.\n\n",
344                "Always returns the full current list."
345            )
346            .to_string(),
347        }
348    }
349
350    pub fn with_description(mut self, description: String) -> Self {
351        self.description = description;
352        self
353    }
354}
355
356impl Tool<BaseCtx> for TodoTool {
357    type Args = TodoArgs;
358    type Output = TodoOutput;
359
360    fn name(&self) -> String {
361        Self::NAME.to_string()
362    }
363
364    fn description(&self) -> String {
365        self.description.clone()
366    }
367
368    fn definition(&self) -> FunctionDefinition {
369        FunctionDefinition {
370            name: self.name(),
371            description: self.description(),
372            parameters: json!({
373                "type": "object",
374                "properties": {
375                    "todos": {
376                        "type": "array",
377                        "description": "Task items to write. Omit to read current list.",
378                        "items": {
379                            "type": "object",
380                            "properties": {
381                                "id": {
382                                    "type": "string",
383                                    "description": "Unique item identifier"
384                                },
385                                "content": {
386                                    "type": "string",
387                                    "description": "Task description"
388                                },
389                                "status": {
390                                    "type": "string",
391                                    "enum": VALID_STATUSES,
392                                    "description": "Current status"
393                                }
394                            },
395                            "required": ["id", "content", "status"],
396                            "additionalProperties": false
397                        }
398                    },
399                    "merge": {
400                        "type": "boolean",
401                        "description": "true: update existing items by id, add new ones. false (default): replace the entire list.",
402                        "default": false
403                    }
404                },
405                "required": [],
406                "additionalProperties": false
407            }),
408            strict: Some(true),
409        }
410    }
411
412    async fn call(
413        &self,
414        ctx: BaseCtx,
415        args: Self::Args,
416        _resources: Vec<Resource>,
417    ) -> Result<ToolOutput<Self::Output>, BoxError> {
418        let hook = ctx.get_state::<TodoToolHook>();
419        let args = if let Some(hook) = &hook {
420            hook.before_tool_call(&ctx, args).await?
421        } else {
422            args
423        };
424
425        let session = todo_session(&ctx);
426        let items = if let Some(todos) = args.todos {
427            session.write(todos, args.merge)
428        } else {
429            session.snapshot()
430        };
431
432        let output = TodoOutput {
433            summary: TodoSummary::from_items(&items),
434            todos: items,
435        };
436
437        if let Some(hook) = &hook {
438            return hook.after_tool_call(&ctx, ToolOutput::new(output)).await;
439        }
440
441        Ok(ToolOutput::new(output))
442    }
443}
444
445fn todo_dedupe_key(item: &TodoItemInput) -> String {
446    let id = item.id.trim();
447    if id.is_empty() {
448        TODO_EMPTY_ID.to_string()
449    } else {
450        id.to_string()
451    }
452}
453
454fn normalize_status(status: Option<&str>) -> String {
455    let status = status
456        .map(|value| value.trim().to_ascii_lowercase())
457        .unwrap_or_else(|| TODO_STATUS_PENDING.to_string());
458
459    if VALID_STATUSES.contains(&status.as_str()) {
460        status
461    } else {
462        TODO_STATUS_PENDING.to_string()
463    }
464}
465
466fn status_marker(status: &str) -> &'static str {
467    match status {
468        TODO_STATUS_PENDING => TODO_MARKER_PENDING,
469        TODO_STATUS_IN_PROGRESS => TODO_MARKER_IN_PROGRESS,
470        TODO_STATUS_COMPLETED => TODO_MARKER_COMPLETED,
471        TODO_STATUS_CANCELLED => TODO_MARKER_CANCELLED,
472        _ => TODO_MARKER_PENDING,
473    }
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479    use crate::engine::EngineBuilder;
480
481    fn input(id: &str, content: Option<&str>, status: Option<&str>) -> TodoItemInput {
482        TodoItemInput {
483            id: id.to_string(),
484            content: content.map(ToString::to_string),
485            status: status.map(ToString::to_string),
486        }
487    }
488
489    fn mock_ctx() -> BaseCtx {
490        EngineBuilder::new().mock_ctx().base
491    }
492
493    #[test]
494    fn replace_mode_dedupes_and_normalizes_items() {
495        let mut store = TodoStore::default();
496
497        let items = store.write(
498            vec![
499                input("1", Some("draft plan"), Some(TODO_STATUS_PENDING)),
500                input("1", Some("final plan"), Some(TODO_STATUS_COMPLETED)),
501                input("", Some(""), Some("invalid")),
502            ],
503            false,
504        );
505
506        assert_eq!(
507            items,
508            vec![
509                TodoItem {
510                    id: "1".to_string(),
511                    content: "final plan".to_string(),
512                    status: TODO_STATUS_COMPLETED.to_string(),
513                },
514                TodoItem {
515                    id: TODO_EMPTY_ID.to_string(),
516                    content: TODO_EMPTY_CONTENT.to_string(),
517                    status: TODO_STATUS_PENDING.to_string(),
518                },
519            ]
520        );
521    }
522
523    #[test]
524    fn merge_mode_updates_existing_items_and_preserves_order() {
525        let mut store = TodoStore::default();
526        store.write(
527            vec![
528                input("1", Some("draft"), Some(TODO_STATUS_PENDING)),
529                input("2", Some("implement"), Some(TODO_STATUS_PENDING)),
530            ],
531            false,
532        );
533
534        let items = store.write(
535            vec![
536                input(
537                    "2",
538                    Some("implement todo tool"),
539                    Some(TODO_STATUS_IN_PROGRESS),
540                ),
541                input("3", Some("write tests"), Some(TODO_STATUS_PENDING)),
542                input("", Some("ignored"), Some(TODO_STATUS_COMPLETED)),
543                input(
544                    "3",
545                    Some("write tests thoroughly"),
546                    Some(TODO_STATUS_COMPLETED),
547                ),
548            ],
549            true,
550        );
551
552        assert_eq!(
553            items,
554            vec![
555                TodoItem {
556                    id: "1".to_string(),
557                    content: "draft".to_string(),
558                    status: TODO_STATUS_PENDING.to_string(),
559                },
560                TodoItem {
561                    id: "2".to_string(),
562                    content: "implement todo tool".to_string(),
563                    status: TODO_STATUS_IN_PROGRESS.to_string(),
564                },
565                TodoItem {
566                    id: "3".to_string(),
567                    content: "write tests thoroughly".to_string(),
568                    status: TODO_STATUS_COMPLETED.to_string(),
569                },
570            ]
571        );
572    }
573
574    #[test]
575    fn injection_format_only_includes_active_items() {
576        let mut store = TodoStore::default();
577        store.write(
578            vec![
579                input("1", Some("plan"), Some(TODO_STATUS_PENDING)),
580                input("2", Some("build"), Some(TODO_STATUS_IN_PROGRESS)),
581                input("3", Some("done"), Some(TODO_STATUS_COMPLETED)),
582                input("4", Some("skip"), Some(TODO_STATUS_CANCELLED)),
583            ],
584            false,
585        );
586
587        let injected = store.format_for_injection().unwrap();
588        assert!(injected.contains(TODO_ACTIVE_LIST_PREFIX));
589        assert!(injected.contains("- [ ] 1. plan (pending)"));
590        assert!(injected.contains("- [>] 2. build (in_progress)"));
591        assert!(!injected.contains("done"));
592        assert!(!injected.contains("skip"));
593    }
594
595    #[tokio::test]
596    async fn tool_call_persists_session_state() {
597        let ctx = mock_ctx();
598        let tool = TodoTool::new();
599
600        let first = tool
601            .call(
602                ctx.clone(),
603                TodoArgs {
604                    todos: Some(vec![input("1", Some("plan"), Some(TODO_STATUS_PENDING))]),
605                    merge: false,
606                },
607                Vec::new(),
608            )
609            .await
610            .unwrap();
611        assert_eq!(first.output.summary.total, 1);
612        assert!(todo_session(&ctx).has_items());
613
614        let second = tool
615            .call(ctx.clone(), TodoArgs::default(), Vec::new())
616            .await
617            .unwrap();
618        assert_eq!(second.output.todos, first.output.todos);
619
620        let third = tool
621            .call(
622                ctx.clone(),
623                TodoArgs {
624                    todos: Some(vec![TodoItemInput {
625                        id: "1".to_string(),
626                        content: Some("plan carefully".to_string()),
627                        status: None,
628                    }]),
629                    merge: true,
630                },
631                Vec::new(),
632            )
633            .await
634            .unwrap();
635
636        assert_eq!(third.output.summary.pending, 1);
637        assert_eq!(third.output.todos[0].content, "plan carefully");
638    }
639}