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::sync::atomic::{AtomicBool, AtomicU64, Ordering};
5use std::sync::Arc;
6use std::time::{Duration, Instant};
7
8use rmcp::model::{DetailedTask, JsonObject, Task, TaskPayload, TaskStatus};
9use serde_json::Value;
10use tokio::sync::Notify;
11use tokio::task::AbortHandle;
12
13use crate::runtime::ShellExecutionToken;
14
15/// Retain completed task results for five minutes unless the client asks for a
16/// different TTL. This is deliberately bounded: Winx task results can contain a
17/// large build log.
18pub const DEFAULT_TASK_TTL_MS: u64 = 5 * 60 * 1_000;
19pub const MAX_TASKS: usize = 32;
20
21#[derive(Debug, Default)]
22pub(crate) struct TaskExecutionControl {
23    cancelled: Arc<AtomicBool>,
24    generation: AtomicU64,
25    execution_token: std::sync::Mutex<Option<ShellExecutionToken>>,
26    launch_finished: AtomicBool,
27    generation_ready: Notify,
28}
29
30impl TaskExecutionControl {
31    fn cancel(&self) {
32        self.cancelled.store(true, Ordering::SeqCst);
33    }
34
35    pub(crate) fn cancellation_flag(&self) -> Arc<AtomicBool> {
36        Arc::clone(&self.cancelled)
37    }
38
39    fn is_cancelled(&self) -> bool {
40        self.cancelled.load(Ordering::SeqCst)
41    }
42
43    fn publish_generation(&self, generation: u64) {
44        self.generation.store(generation, Ordering::SeqCst);
45        self.generation_ready.notify_waiters();
46    }
47
48    fn publish_execution(&self, token: ShellExecutionToken) {
49        let generation = token.generation;
50        if let Ok(mut slot) = self.execution_token.lock() {
51            *slot = Some(token);
52        }
53        self.publish_generation(generation);
54    }
55
56    fn execution_token(&self) -> Option<ShellExecutionToken> {
57        self.execution_token.lock().ok().and_then(|slot| slot.clone())
58    }
59
60    fn generation(&self) -> Option<u64> {
61        match self.generation.load(Ordering::SeqCst) {
62            0 => None,
63            generation => Some(generation),
64        }
65    }
66
67    pub(crate) fn finish_launch(&self) {
68        self.launch_finished.store(true, Ordering::SeqCst);
69        self.generation_ready.notify_waiters();
70    }
71
72    pub(crate) async fn wait_for_generation(&self) -> Option<u64> {
73        loop {
74            let notified = self.generation_ready.notified();
75            if let Some(generation) = self.generation() {
76                return Some(generation);
77            }
78            if self.launch_finished.load(Ordering::SeqCst) {
79                return None;
80            }
81            notified.await;
82        }
83    }
84
85    pub(crate) async fn wait_for_execution(&self) -> Option<ShellExecutionToken> {
86        self.wait_for_generation().await?;
87        self.execution_token()
88    }
89}
90
91#[derive(Debug)]
92pub struct TaskEntry {
93    pub task: Task,
94    pub result: Option<Value>,
95    pub abort_handle: Option<AbortHandle>,
96    pub thread_id: String,
97    command_generation: Option<u64>,
98    execution_token: Option<ShellExecutionToken>,
99    execution_control: Arc<TaskExecutionControl>,
100    expires_at: Option<Instant>,
101    terminal_at: Option<Instant>,
102}
103
104impl TaskEntry {
105    pub fn working(task: Task, thread_id: String) -> Self {
106        Self {
107            task,
108            result: None,
109            abort_handle: None,
110            thread_id,
111            command_generation: None,
112            execution_token: None,
113            execution_control: Arc::new(TaskExecutionControl::default()),
114            expires_at: None,
115            terminal_at: None,
116        }
117    }
118
119    pub fn command_generation(&self) -> Option<u64> {
120        self.command_generation.or_else(|| self.execution_control.generation())
121    }
122
123    pub fn set_command_generation(&mut self, generation: u64) {
124        self.execution_control.publish_generation(generation);
125        self.command_generation = Some(generation);
126    }
127
128    pub fn execution_token(&self) -> Option<ShellExecutionToken> {
129        self.execution_token.clone().or_else(|| self.execution_control.execution_token())
130    }
131
132    pub fn set_execution_token(&mut self, token: ShellExecutionToken) {
133        self.command_generation = Some(token.generation);
134        self.execution_control.publish_execution(token.clone());
135        self.execution_token = Some(token);
136    }
137
138    pub fn request_cancel(&self) {
139        self.execution_control.cancel();
140    }
141
142    pub fn is_cancel_requested(&self) -> bool {
143        self.execution_control.is_cancelled()
144    }
145
146    pub(crate) fn execution_control(&self) -> Arc<TaskExecutionControl> {
147        Arc::clone(&self.execution_control)
148    }
149
150    pub fn finish(
151        &mut self,
152        status: TaskStatus,
153        status_message: Option<String>,
154        result: Option<Value>,
155    ) {
156        self.task.status = status;
157        self.task.status_message = status_message;
158        self.task.last_updated_at = rmcp::task_manager::current_timestamp();
159        self.result = result;
160        self.abort_handle = None;
161        self.expires_at = self.task.ttl_ms.map(|ttl| Instant::now() + Duration::from_millis(ttl));
162        self.terminal_at = Some(Instant::now());
163    }
164
165    pub fn detailed(&self) -> DetailedTask {
166        let payload = match self.task.status {
167            TaskStatus::Working => TaskPayload::Working,
168            TaskStatus::InputRequired => TaskPayload::Failed {
169                error: task_error("Winx BashCommand tasks do not request in-task input"),
170            },
171            TaskStatus::Completed => match self.result.clone() {
172                Some(Value::Object(result)) => TaskPayload::Completed { result },
173                _ => {
174                    TaskPayload::Failed { error: task_error("Completed task has no object result") }
175                }
176            },
177            TaskStatus::Failed => TaskPayload::Failed {
178                error: task_error(
179                    self.task
180                        .status_message
181                        .as_deref()
182                        .unwrap_or("Task failed without a status message"),
183                ),
184            },
185            TaskStatus::Cancelled => TaskPayload::Cancelled,
186            _ => TaskPayload::Failed { error: task_error("Task has an unsupported status") },
187        };
188        DetailedTask::new(self.task.clone(), payload)
189    }
190
191    fn is_expired(&self, now: Instant) -> bool {
192        self.expires_at.is_some_and(|expiry| now >= expiry)
193    }
194}
195
196#[derive(Debug, Default)]
197pub struct TaskRegistry {
198    entries: HashMap<String, TaskEntry>,
199}
200
201impl TaskRegistry {
202    pub fn prune(&mut self) {
203        let now = Instant::now();
204        self.entries.retain(|_, entry| !entry.is_expired(now));
205    }
206
207    pub fn insert(&mut self, task_id: String, entry: TaskEntry) -> Result<(), &'static str> {
208        self.prune();
209        while self.entries.len() >= MAX_TASKS {
210            let oldest_terminal = self
211                .entries
212                .iter()
213                .filter_map(|(id, entry)| entry.terminal_at.map(|at| (id.clone(), at)))
214                .min_by_key(|(_, at)| *at)
215                .map(|(id, _)| id);
216            let Some(oldest_terminal) = oldest_terminal else { break };
217            self.entries.remove(&oldest_terminal);
218        }
219        let working = self
220            .entries
221            .values()
222            .filter(|entry| {
223                matches!(entry.task.status, TaskStatus::Working | TaskStatus::InputRequired)
224            })
225            .count();
226        if working >= MAX_TASKS {
227            return Err("MCP task limit reached; wait for or cancel an existing task");
228        }
229        if self.entries.contains_key(&task_id) {
230            return Err("MCP task identifier collision");
231        }
232        self.entries.insert(task_id, entry);
233        Ok(())
234    }
235
236    pub fn get(&mut self, task_id: &str) -> Option<&TaskEntry> {
237        self.prune();
238        self.entries.get(task_id)
239    }
240
241    pub fn get_mut(&mut self, task_id: &str) -> Option<&mut TaskEntry> {
242        self.prune();
243        self.entries.get_mut(task_id)
244    }
245
246    pub fn remove(&mut self, task_id: &str) -> Option<TaskEntry> {
247        self.entries.remove(task_id)
248    }
249}
250
251fn task_error(message: &str) -> JsonObject {
252    serde_json::Map::from_iter([
253        ("code".to_string(), Value::from(-32603)),
254        ("message".to_string(), Value::from(message)),
255    ])
256}
257
258#[cfg(test)]
259mod tests {
260    #![allow(clippy::unwrap_used)]
261
262    use super::*;
263
264    fn task(id: &str, ttl: u64) -> Task {
265        Task::new(
266            id.to_string(),
267            TaskStatus::Working,
268            "2026-07-21T12:00:00Z".to_string(),
269            "2026-07-21T12:00:00Z".to_string(),
270        )
271        .with_ttl_ms(ttl)
272    }
273
274    #[test]
275    fn terminal_task_keeps_repeatable_result() {
276        let mut registry = TaskRegistry::default();
277        registry
278            .insert("one".into(), TaskEntry::working(task("one", 60_000), "tid_1".into()))
279            .unwrap();
280        registry.get_mut("one").unwrap().finish(
281            TaskStatus::Completed,
282            Some("done".into()),
283            Some(serde_json::json!({"content": []})),
284        );
285
286        assert_eq!(
287            registry.get("one").and_then(|entry| entry.result.clone()),
288            Some(serde_json::json!({"content": []}))
289        );
290        assert_eq!(
291            registry.get("one").and_then(|entry| entry.result.clone()),
292            Some(serde_json::json!({"content": []}))
293        );
294    }
295
296    #[test]
297    fn terminal_tasks_are_evicted_before_they_exhaust_working_capacity() {
298        let mut registry = TaskRegistry::default();
299        for index in 0..MAX_TASKS {
300            let id = format!("terminal-{index}");
301            registry
302                .insert(id.clone(), TaskEntry::working(task(&id, 300_000), "tid".into()))
303                .unwrap();
304            registry.get_mut(&id).unwrap().finish(
305                TaskStatus::Completed,
306                Some("done".into()),
307                Some(serde_json::json!({"content": []})),
308            );
309        }
310
311        registry
312            .insert(
313                "new-working".into(),
314                TaskEntry::working(task("new-working", 300_000), "tid".into()),
315            )
316            .unwrap();
317        assert_eq!(registry.get("new-working").unwrap().task.status, TaskStatus::Working);
318    }
319}
winx_code_agent/state/task_state.rs

winx_code_agent/state/
task_state.rs
