walrus_daemon/hook/system/task/
mod.rs

1//! Lightweight task set — tracks delegated agent work as JoinHandles.
2//!
3//! Replaces the old TaskRegistry. A task is either running or done.
4//! No status state machine, no queuing, no approval inbox.
5
6use std::{
7    collections::HashMap,
8    sync::atomic::{AtomicU64, Ordering},
9};
10use tokio::task::JoinHandle;
11use tokio::time::Instant;
12use wcore::protocol::message::TaskInfo;
13
14pub(crate) mod tool;
15
16/// A delegated unit of agent work.
17pub struct Task {
18    /// Unique task identifier.
19    pub id: u64,
20    /// Agent assigned to this task.
21    pub agent: String,
22    /// Human-readable task description / message.
23    pub description: String,
24    /// When this task was created.
25    pub created_at: Instant,
26    /// Session allocated for this task's execution.
27    pub session_id: Option<u64>,
28    /// Background handle — resolves to the agent's final response.
29    pub handle: Option<JoinHandle<String>>,
30    /// Cached result after handle resolves.
31    pub result: Option<String>,
32    /// Cached error after handle resolves.
33    pub error: Option<String>,
34}
35
36impl Task {
37    /// Whether the task has completed (handle resolved or result cached).
38    pub fn is_done(&self) -> bool {
39        self.result.is_some()
40            || self.error.is_some()
41            || self.handle.as_ref().is_some_and(|h| h.is_finished())
42    }
43
44    /// Status string for protocol compatibility.
45    pub fn status(&self) -> &'static str {
46        if self.error.is_some() {
47            "failed"
48        } else if self.result.is_some() || self.handle.as_ref().is_some_and(|h| h.is_finished()) {
49            "finished"
50        } else {
51            "in_progress"
52        }
53    }
54
55    /// Build a `TaskInfo` snapshot for the wire protocol.
56    pub fn to_info(&self) -> TaskInfo {
57        TaskInfo {
58            id: self.id,
59            parent_id: None,
60            agent: self.agent.to_string(),
61            status: self.status().to_string(),
62            description: self.description.clone(),
63            result: self.result.clone(),
64            error: self.error.clone(),
65            created_by: String::new(),
66            prompt_tokens: 0,
67            completion_tokens: 0,
68            alive_secs: self.created_at.elapsed().as_secs(),
69            blocked_on: None,
70        }
71    }
72}
73
74/// Lightweight task tracker — just a map of JoinHandles with metadata.
75#[derive(Default)]
76pub struct TaskSet {
77    tasks: HashMap<u64, Task>,
78    next_id: AtomicU64,
79}
80
81impl TaskSet {
82    pub fn new() -> Self {
83        Self {
84            tasks: HashMap::new(),
85            next_id: AtomicU64::new(1),
86        }
87    }
88
89    /// Insert a new task. Returns the task ID.
90    pub fn insert(&mut self, agent: String, description: String) -> u64 {
91        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
92        self.tasks.insert(
93            id,
94            Task {
95                id,
96                agent,
97                description,
98                created_at: Instant::now(),
99                session_id: None,
100                handle: None,
101                result: None,
102                error: None,
103            },
104        );
105        id
106    }
107
108    /// Get a reference to a task.
109    pub fn get(&self, id: u64) -> Option<&Task> {
110        self.tasks.get(&id)
111    }
112
113    /// Get a mutable reference to a task.
114    pub fn get_mut(&mut self, id: u64) -> Option<&mut Task> {
115        self.tasks.get_mut(&id)
116    }
117
118    /// List all tasks (most recent first), up to `limit`.
119    pub fn list(&self, limit: usize) -> Vec<&Task> {
120        let mut tasks: Vec<_> = self.tasks.values().collect();
121        tasks.sort_by(|a, b| b.id.cmp(&a.id));
122        tasks.truncate(limit);
123        tasks
124    }
125
126    /// Remove a task by ID, aborting its handle if running.
127    pub fn kill(&mut self, id: u64) -> bool {
128        if let Some(task) = self.tasks.remove(&id) {
129            if let Some(handle) = task.handle {
130                handle.abort();
131            }
132            true
133        } else {
134            false
135        }
136    }
137}
walrus_daemon/hook/system/task/mod.rs

walrus_daemon/hook/system/task/
mod.rs
