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, or with todos=null in strict schemas, ",
335                "to read the current list.\n\n",
336                "Writing:\n",
337                "- Provide 'todos' array to create/update items\n",
338                "- merge=false (default): replace the entire list with a fresh plan\n",
339                "- merge=true: update existing items by id, add any new ones\n\n",
340                "Each item: {id: string, content: string, ",
341                "status: pending|in_progress|completed|cancelled}\n",
342                "List order is priority. Only ONE item in_progress at a time.\n",
343                "Mark items completed immediately when done. If something fails, ",
344                "cancel it and add a revised item.\n\n",
345                "Always returns the full current list."
346            )
347            .to_string(),
348        }
349    }
350
351    pub fn with_description(mut self, description: String) -> Self {
352        self.description = description;
353        self
354    }
355}
356
357impl Tool<BaseCtx> for TodoTool {
358    type Args = TodoArgs;
359    type Output = TodoOutput;
360
361    fn name(&self) -> String {
362        Self::NAME.to_string()
363    }
364
365    fn description(&self) -> String {
366        self.description.clone()
367    }
368
369    fn definition(&self) -> FunctionDefinition {
370        FunctionDefinition {
371            name: self.name(),
372            description: self.description(),
373            parameters: json!({
374                "type": "object",
375                "properties": {
376                    "todos": {
377                        "description": "Task items to write. Use null to read current list.",
378                        "type": ["array", "null"],
379                        "items": {
380                            "type": "object",
381                            "properties": {
382                                "id": {
383                                    "type": "string",
384                                    "description": "Unique item identifier"
385                                },
386                                "content": {
387                                    "type": ["string", "null"],
388                                    "description": "Task description. Use null to leave it unchanged when merge=true."
389                                },
390                                "status": {
391                                    "type": ["string", "null"],
392                                    "enum": [
393                                        TODO_STATUS_PENDING,
394                                        TODO_STATUS_IN_PROGRESS,
395                                        TODO_STATUS_COMPLETED,
396                                        TODO_STATUS_CANCELLED,
397                                        null
398                                    ],
399                                    "description": "Current status. Use null to keep the existing status when merge=true."
400                                }
401                            },
402                            "required": ["id", "content", "status"],
403                            "additionalProperties": false
404                        }
405                    },
406                    "merge": {
407                        "type": "boolean",
408                        "description": "true: update existing items by id, add new ones. false (default): replace the entire list.",
409                        "default": false
410                    }
411                },
412                "required": ["todos", "merge"],
413                "additionalProperties": false
414            }),
415            strict: Some(true),
416        }
417    }
418
419    async fn call(
420        &self,
421        ctx: BaseCtx,
422        args: Self::Args,
423        _resources: Vec<Resource>,
424    ) -> Result<ToolOutput<Self::Output>, BoxError> {
425        let hook = ctx.get_state::<TodoToolHook>();
426        let args = if let Some(hook) = &hook {
427            hook.before_tool_call(&ctx, args).await?
428        } else {
429            args
430        };
431
432        let session = todo_session(&ctx);
433        let items = if let Some(todos) = args.todos {
434            session.write(todos, args.merge)
435        } else {
436            session.snapshot()
437        };
438
439        let output = TodoOutput {
440            summary: TodoSummary::from_items(&items),
441            todos: items,
442        };
443
444        if let Some(hook) = &hook {
445            return hook.after_tool_call(&ctx, ToolOutput::new(output)).await;
446        }
447
448        Ok(ToolOutput::new(output))
449    }
450}
451
452fn todo_dedupe_key(item: &TodoItemInput) -> String {
453    let id = item.id.trim();
454    if id.is_empty() {
455        TODO_EMPTY_ID.to_string()
456    } else {
457        id.to_string()
458    }
459}
460
461fn normalize_status(status: Option<&str>) -> String {
462    let status = status
463        .map(|value| value.trim().to_ascii_lowercase())
464        .unwrap_or_else(|| TODO_STATUS_PENDING.to_string());
465
466    if VALID_STATUSES.contains(&status.as_str()) {
467        status
468    } else {
469        TODO_STATUS_PENDING.to_string()
470    }
471}
472
473fn status_marker(status: &str) -> &'static str {
474    match status {
475        TODO_STATUS_PENDING => TODO_MARKER_PENDING,
476        TODO_STATUS_IN_PROGRESS => TODO_MARKER_IN_PROGRESS,
477        TODO_STATUS_COMPLETED => TODO_MARKER_COMPLETED,
478        TODO_STATUS_CANCELLED => TODO_MARKER_CANCELLED,
479        _ => TODO_MARKER_PENDING,
480    }
481}
482
483#[cfg(test)]
484mod tests {
485    use super::*;
486    use crate::engine::EngineBuilder;
487
488    fn input(id: &str, content: Option<&str>, status: Option<&str>) -> TodoItemInput {
489        TodoItemInput {
490            id: id.to_string(),
491            content: content.map(ToString::to_string),
492            status: status.map(ToString::to_string),
493        }
494    }
495
496    fn mock_ctx() -> BaseCtx {
497        EngineBuilder::new().mock_ctx().base
498    }
499
500    #[test]
501    fn replace_mode_dedupes_and_normalizes_items() {
502        let mut store = TodoStore::default();
503
504        let items = store.write(
505            vec![
506                input("1", Some("draft plan"), Some(TODO_STATUS_PENDING)),
507                input("1", Some("final plan"), Some(TODO_STATUS_COMPLETED)),
508                input("", Some(""), Some("invalid")),
509            ],
510            false,
511        );
512
513        assert_eq!(
514            items,
515            vec![
516                TodoItem {
517                    id: "1".to_string(),
518                    content: "final plan".to_string(),
519                    status: TODO_STATUS_COMPLETED.to_string(),
520                },
521                TodoItem {
522                    id: TODO_EMPTY_ID.to_string(),
523                    content: TODO_EMPTY_CONTENT.to_string(),
524                    status: TODO_STATUS_PENDING.to_string(),
525                },
526            ]
527        );
528    }
529
530    #[test]
531    fn merge_mode_updates_existing_items_and_preserves_order() {
532        let mut store = TodoStore::default();
533        store.write(
534            vec![
535                input("1", Some("draft"), Some(TODO_STATUS_PENDING)),
536                input("2", Some("implement"), Some(TODO_STATUS_PENDING)),
537            ],
538            false,
539        );
540
541        let items = store.write(
542            vec![
543                input(
544                    "2",
545                    Some("implement todo tool"),
546                    Some(TODO_STATUS_IN_PROGRESS),
547                ),
548                input("3", Some("write tests"), Some(TODO_STATUS_PENDING)),
549                input("", Some("ignored"), Some(TODO_STATUS_COMPLETED)),
550                input(
551                    "3",
552                    Some("write tests thoroughly"),
553                    Some(TODO_STATUS_COMPLETED),
554                ),
555            ],
556            true,
557        );
558
559        assert_eq!(
560            items,
561            vec![
562                TodoItem {
563                    id: "1".to_string(),
564                    content: "draft".to_string(),
565                    status: TODO_STATUS_PENDING.to_string(),
566                },
567                TodoItem {
568                    id: "2".to_string(),
569                    content: "implement todo tool".to_string(),
570                    status: TODO_STATUS_IN_PROGRESS.to_string(),
571                },
572                TodoItem {
573                    id: "3".to_string(),
574                    content: "write tests thoroughly".to_string(),
575                    status: TODO_STATUS_COMPLETED.to_string(),
576                },
577            ]
578        );
579    }
580
581    #[test]
582    fn injection_format_only_includes_active_items() {
583        let mut store = TodoStore::default();
584        store.write(
585            vec![
586                input("1", Some("plan"), Some(TODO_STATUS_PENDING)),
587                input("2", Some("build"), Some(TODO_STATUS_IN_PROGRESS)),
588                input("3", Some("done"), Some(TODO_STATUS_COMPLETED)),
589                input("4", Some("skip"), Some(TODO_STATUS_CANCELLED)),
590            ],
591            false,
592        );
593
594        let injected = store.format_for_injection().unwrap();
595        assert!(injected.contains(TODO_ACTIVE_LIST_PREFIX));
596        assert!(injected.contains("- [ ] 1. plan (pending)"));
597        assert!(injected.contains("- [>] 2. build (in_progress)"));
598        assert!(!injected.contains("done"));
599        assert!(!injected.contains("skip"));
600    }
601
602    #[tokio::test]
603    async fn tool_call_persists_session_state() {
604        let ctx = mock_ctx();
605        let tool = TodoTool::new();
606
607        let first = tool
608            .call(
609                ctx.clone(),
610                TodoArgs {
611                    todos: Some(vec![input("1", Some("plan"), Some(TODO_STATUS_PENDING))]),
612                    merge: false,
613                },
614                Vec::new(),
615            )
616            .await
617            .unwrap();
618        assert_eq!(first.output.summary.total, 1);
619        assert!(todo_session(&ctx).has_items());
620
621        let second = tool
622            .call(ctx.clone(), TodoArgs::default(), Vec::new())
623            .await
624            .unwrap();
625        assert_eq!(second.output.todos, first.output.todos);
626
627        let third = tool
628            .call(
629                ctx.clone(),
630                TodoArgs {
631                    todos: Some(vec![TodoItemInput {
632                        id: "1".to_string(),
633                        content: Some("plan carefully".to_string()),
634                        status: None,
635                    }]),
636                    merge: true,
637                },
638                Vec::new(),
639            )
640            .await
641            .unwrap();
642
643        assert_eq!(third.output.summary.pending, 1);
644        assert_eq!(third.output.todos[0].content, "plan carefully");
645    }
646
647    #[test]
648    fn definition_schema_avoids_anyof() {
649        let definition = TodoTool::new().definition();
650
651        assert!(
652            definition.parameters["properties"]["todos"]
653                .get("anyOf")
654                .is_none()
655        );
656        assert_eq!(
657            definition.parameters["properties"]["todos"]["type"],
658            json!(["array", "null"])
659        );
660        assert_eq!(
661            definition.parameters["properties"]["todos"]["items"]["properties"]["content"]["type"],
662            json!(["string", "null"])
663        );
664        assert_eq!(
665            definition.parameters["properties"]["todos"]["items"]["properties"]["status"]["enum"],
666            json!([
667                TODO_STATUS_PENDING,
668                TODO_STATUS_IN_PROGRESS,
669                TODO_STATUS_COMPLETED,
670                TODO_STATUS_CANCELLED,
671                null
672            ])
673        );
674        assert_eq!(definition.parameters["required"], json!(["todos", "merge"]));
675    }
676}