winx_code_agent/state/
task_state.rs

1//! In-memory lifecycle state for the SEP-2663 MCP Tasks extension.
2
3use std::collections::HashMap;
4use std::time::{Duration, Instant};
5
6use rmcp::model::{DetailedTask, JsonObject, Task, TaskPayload, 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 MAX_TASKS: usize = 32;
15
16#[derive(Debug)]
17pub struct TaskEntry {
18    pub task: Task,
19    pub result: Option<Value>,
20    pub abort_handle: Option<AbortHandle>,
21    pub thread_id: String,
22    expires_at: Option<Instant>,
23}
24
25impl TaskEntry {
26    pub fn working(task: Task, thread_id: String) -> Self {
27        Self { task, result: None, abort_handle: None, thread_id, expires_at: None }
28    }
29
30    pub fn finish(
31        &mut self,
32        status: TaskStatus,
33        status_message: Option<String>,
34        result: Option<Value>,
35    ) {
36        self.task.status = status;
37        self.task.status_message = status_message;
38        self.task.last_updated_at = rmcp::task_manager::current_timestamp();
39        self.result = result;
40        self.abort_handle = None;
41        self.expires_at = self.task.ttl_ms.map(|ttl| Instant::now() + Duration::from_millis(ttl));
42    }
43
44    pub fn detailed(&self) -> DetailedTask {
45        let payload = match self.task.status {
46            TaskStatus::Working => TaskPayload::Working,
47            TaskStatus::InputRequired => TaskPayload::Failed {
48                error: task_error("Winx BashCommand tasks do not request in-task input"),
49            },
50            TaskStatus::Completed => match self.result.clone() {
51                Some(Value::Object(result)) => TaskPayload::Completed { result },
52                _ => {
53                    TaskPayload::Failed { error: task_error("Completed task has no object result") }
54                }
55            },
56            TaskStatus::Failed => TaskPayload::Failed {
57                error: task_error(
58                    self.task
59                        .status_message
60                        .as_deref()
61                        .unwrap_or("Task failed without a status message"),
62                ),
63            },
64            TaskStatus::Cancelled => TaskPayload::Cancelled,
65            _ => TaskPayload::Failed { error: task_error("Task has an unsupported status") },
66        };
67        DetailedTask::new(self.task.clone(), payload)
68    }
69
70    fn is_expired(&self, now: Instant) -> bool {
71        self.expires_at.is_some_and(|expiry| now >= expiry)
72    }
73}
74
75#[derive(Debug, Default)]
76pub struct TaskRegistry {
77    entries: HashMap<String, TaskEntry>,
78}
79
80impl TaskRegistry {
81    pub fn prune(&mut self) {
82        let now = Instant::now();
83        self.entries.retain(|_, entry| !entry.is_expired(now));
84    }
85
86    pub fn insert(&mut self, task_id: String, entry: TaskEntry) -> Result<(), &'static str> {
87        self.prune();
88        if self.entries.len() >= MAX_TASKS {
89            return Err("MCP task limit reached; wait for or cancel an existing task");
90        }
91        if self.entries.contains_key(&task_id) {
92            return Err("MCP task identifier collision");
93        }
94        self.entries.insert(task_id, entry);
95        Ok(())
96    }
97
98    pub fn get(&mut self, task_id: &str) -> Option<&TaskEntry> {
99        self.prune();
100        self.entries.get(task_id)
101    }
102
103    pub fn get_mut(&mut self, task_id: &str) -> Option<&mut TaskEntry> {
104        self.prune();
105        self.entries.get_mut(task_id)
106    }
107}
108
109fn task_error(message: &str) -> JsonObject {
110    serde_json::Map::from_iter([
111        ("code".to_string(), Value::from(-32603)),
112        ("message".to_string(), Value::from(message)),
113    ])
114}
115
116#[cfg(test)]
117mod tests {
118    #![allow(clippy::unwrap_used)]
119
120    use super::*;
121
122    fn task(id: &str, ttl: u64) -> Task {
123        Task::new(
124            id.to_string(),
125            TaskStatus::Working,
126            "2026-07-21T12:00:00Z".to_string(),
127            "2026-07-21T12:00:00Z".to_string(),
128        )
129        .with_ttl_ms(ttl)
130    }
131
132    #[test]
133    fn terminal_task_keeps_repeatable_result() {
134        let mut registry = TaskRegistry::default();
135        registry
136            .insert("one".into(), TaskEntry::working(task("one", 60_000), "tid_1".into()))
137            .unwrap();
138        registry.get_mut("one").unwrap().finish(
139            TaskStatus::Completed,
140            Some("done".into()),
141            Some(serde_json::json!({"content": []})),
142        );
143
144        assert_eq!(
145            registry.get("one").and_then(|entry| entry.result.clone()),
146            Some(serde_json::json!({"content": []}))
147        );
148        assert_eq!(
149            registry.get("one").and_then(|entry| entry.result.clone()),
150            Some(serde_json::json!({"content": []}))
151        );
152    }
153}
winx_code_agent/state/task_state.rs

winx_code_agent/state/
task_state.rs
