winx_code_agent/state/
task_state.rs

1//! In-memory lifecycle state for MCP Tasks.
2
3use std::collections::HashMap;
4use std::time::{Duration, Instant};
5
6use rmcp::model::{Task, TaskStatus};
7use serde_json::Value;
8use tokio::task::AbortHandle;
9
10/// Retain completed task results for five minutes unless the client asks for a
11/// different TTL. This is deliberately bounded: Winx task results can contain a
12/// large build log.
13pub const DEFAULT_TASK_TTL_MS: u64 = 5 * 60 * 1_000;
14pub const MIN_TASK_TTL_MS: u64 = 1_000;
15pub const MAX_TASK_TTL_MS: u64 = 24 * 60 * 60 * 1_000;
16pub const MAX_TASKS: usize = 32;
17
18#[derive(Debug)]
19pub struct TaskEntry {
20    pub task: Task,
21    pub result: Option<Value>,
22    pub abort_handle: Option<AbortHandle>,
23    pub thread_id: String,
24    expires_at: Option<Instant>,
25}
26
27impl TaskEntry {
28    pub fn working(task: Task, thread_id: String) -> Self {
29        Self { task, result: None, abort_handle: None, thread_id, expires_at: None }
30    }
31
32    pub fn finish(
33        &mut self,
34        status: TaskStatus,
35        status_message: Option<String>,
36        result: Option<Value>,
37    ) {
38        self.task.status = status;
39        self.task.status_message = status_message;
40        self.task.last_updated_at = rmcp::task_manager::current_timestamp();
41        self.result = result;
42        self.abort_handle = None;
43        self.expires_at = self.task.ttl.map(|ttl| Instant::now() + Duration::from_millis(ttl));
44    }
45
46    fn is_expired(&self, now: Instant) -> bool {
47        self.expires_at.is_some_and(|expiry| now >= expiry)
48    }
49}
50
51#[derive(Debug, Default)]
52pub struct TaskRegistry {
53    entries: HashMap<String, TaskEntry>,
54}
55
56impl TaskRegistry {
57    pub fn prune(&mut self) {
58        let now = Instant::now();
59        self.entries.retain(|_, entry| !entry.is_expired(now));
60    }
61
62    pub fn insert(&mut self, task_id: String, entry: TaskEntry) -> Result<(), &'static str> {
63        self.prune();
64        if self.entries.len() >= MAX_TASKS {
65            return Err("MCP task limit reached; wait for or cancel an existing task");
66        }
67        if self.entries.contains_key(&task_id) {
68            return Err("MCP task identifier collision");
69        }
70        self.entries.insert(task_id, entry);
71        Ok(())
72    }
73
74    pub fn get(&mut self, task_id: &str) -> Option<&TaskEntry> {
75        self.prune();
76        self.entries.get(task_id)
77    }
78
79    pub fn get_mut(&mut self, task_id: &str) -> Option<&mut TaskEntry> {
80        self.prune();
81        self.entries.get_mut(task_id)
82    }
83
84    pub fn list(&mut self) -> Vec<Task> {
85        self.prune();
86        let mut tasks = self.entries.values().map(|entry| entry.task.clone()).collect::<Vec<_>>();
87        tasks.sort_by(|a, b| b.created_at.cmp(&a.created_at));
88        tasks
89    }
90}
91
92/// Negotiate a bounded result-retention TTL from the task augmentation object.
93pub fn requested_ttl(task: Option<&serde_json::Map<String, Value>>) -> u64 {
94    task.and_then(|metadata| metadata.get("ttl"))
95        .and_then(Value::as_u64)
96        .unwrap_or(DEFAULT_TASK_TTL_MS)
97        .clamp(MIN_TASK_TTL_MS, MAX_TASK_TTL_MS)
98}
99
100#[cfg(test)]
101mod tests {
102    #![allow(clippy::unwrap_used)]
103
104    use super::*;
105
106    fn task(id: &str, ttl: u64) -> Task {
107        Task::new(
108            id.to_string(),
109            TaskStatus::Working,
110            "2026-07-21T12:00:00Z".to_string(),
111            "2026-07-21T12:00:00Z".to_string(),
112        )
113        .with_ttl(ttl)
114    }
115
116    #[test]
117    fn requested_ttl_is_bounded() {
118        let low = serde_json::from_value(serde_json::json!({"ttl": 1})).unwrap();
119        let high = serde_json::from_value(serde_json::json!({"ttl": u64::MAX})).unwrap();
120        assert_eq!(requested_ttl(Some(&low)), MIN_TASK_TTL_MS);
121        assert_eq!(requested_ttl(Some(&high)), MAX_TASK_TTL_MS);
122        assert_eq!(requested_ttl(None), DEFAULT_TASK_TTL_MS);
123    }
124
125    #[test]
126    fn terminal_task_keeps_repeatable_result() {
127        let mut registry = TaskRegistry::default();
128        registry
129            .insert("one".into(), TaskEntry::working(task("one", 60_000), "tid_1".into()))
130            .unwrap();
131        registry.get_mut("one").unwrap().finish(
132            TaskStatus::Completed,
133            Some("done".into()),
134            Some(serde_json::json!({"content": []})),
135        );
136
137        assert_eq!(
138            registry.get("one").and_then(|entry| entry.result.clone()),
139            Some(serde_json::json!({"content": []}))
140        );
141        assert_eq!(
142            registry.get("one").and_then(|entry| entry.result.clone()),
143            Some(serde_json::json!({"content": []}))
144        );
145    }
146}
winx_code_agent/state/task_state.rs

winx_code_agent/state/
task_state.rs
