projectatlas/mcp/
task_registry.rs

1//! Purpose: Own the bounded session-local MCP task registry and task lifecycle values.
2
3use super::{
4    MCP_TASK_CONTRACT_ID, MCP_TASK_PROGRESS_CONTRACT_MESSAGE, MCP_TASK_REGISTRY_CAPACITY,
5    MCP_TOOL_ATLAS_TASK_STATUS, mcp_unix_time_ms,
6};
7use projectatlas_core::IndexWorkControl;
8use serde::Serialize;
9use std::collections::VecDeque;
10
11/// Bounded in-memory registry for MCP task-progress records.
12#[derive(Debug, Clone)]
13pub(super) struct McpTaskRegistry {
14    /// Session-local task records.
15    records: VecDeque<McpTaskRecord>,
16}
17
18impl McpTaskRegistry {
19    /// Create a registry with the built-in task-progress contract record.
20    pub(super) fn new() -> Self {
21        let now = mcp_unix_time_ms();
22        let mut registry = Self {
23            records: VecDeque::new(),
24        };
25        registry.insert(McpTaskRecord {
26            task_id: MCP_TASK_CONTRACT_ID.to_string(),
27            operation: McpTaskOperation::Contract,
28            state: McpTaskState::Complete,
29            created_at_ms: now,
30            updated_at_ms: now,
31            progress: Some(McpTaskProgress {
32                current: Some(1),
33                total: Some(1),
34                message: Some(MCP_TASK_PROGRESS_CONTRACT_MESSAGE.to_string()),
35            }),
36            error: None,
37            result_ref: Some(MCP_TOOL_ATLAS_TASK_STATUS.to_string()),
38            cancelable: false,
39            control: None,
40        });
41        registry
42    }
43
44    /// Return the number of admitted tasks that have not reached a terminal state.
45    pub(super) fn active_count(&self) -> usize {
46        self.records
47            .iter()
48            .filter(|record| !record.is_terminal_state())
49            .count()
50    }
51
52    /// Return the number of retained task records.
53    #[cfg(test)]
54    pub(super) fn len(&self) -> usize {
55        self.records.len()
56    }
57
58    /// Return the newest retained task id for an operation.
59    #[cfg(test)]
60    pub(super) fn latest_task_id(&self, operation: &McpTaskOperation) -> Option<String> {
61        self.records
62            .iter()
63            .rev()
64            .find(|record| &record.operation == operation)
65            .map(|record| record.task_id.clone())
66    }
67
68    /// Insert or replace one task record while preserving the fixed registry capacity.
69    pub(super) fn insert(&mut self, record: McpTaskRecord) {
70        if let Some(existing_index) = self
71            .records
72            .iter()
73            .position(|current| current.task_id == record.task_id)
74        {
75            let _removed = self.records.remove(existing_index);
76        }
77        while self.records.len() >= MCP_TASK_REGISTRY_CAPACITY {
78            if let Some(finished_index) = self
79                .records
80                .iter()
81                .position(McpTaskRecord::is_terminal_state)
82            {
83                let _evicted = self.records.remove(finished_index);
84            } else {
85                let _evicted = self.records.pop_front();
86            }
87        }
88        self.records.push_back(record);
89    }
90
91    /// Return a task record by id.
92    pub(super) fn get(&self, task_id: &str) -> Option<McpTaskRecord> {
93        self.records
94            .iter()
95            .find(|record| record.task_id == task_id)
96            .cloned()
97    }
98
99    /// Update a matching task through a bounded mutable pass.
100    pub(super) fn update<F>(&mut self, task_id: &str, update: F) -> Option<McpTaskRecord>
101    where
102        F: FnOnce(&mut McpTaskRecord),
103    {
104        let record = self
105            .records
106            .iter_mut()
107            .find(|record| record.task_id == task_id)?;
108        update(record);
109        Some(record.clone())
110    }
111}
112
113/// One MCP task-progress record.
114#[derive(Debug, Clone, Serialize)]
115pub(super) struct McpTaskRecord {
116    /// Opaque session-local task id.
117    pub(super) task_id: String,
118    /// Operation family.
119    pub(super) operation: McpTaskOperation,
120    /// Current task state.
121    pub(super) state: McpTaskState,
122    /// Creation timestamp in Unix milliseconds.
123    pub(super) created_at_ms: u128,
124    /// Last update timestamp in Unix milliseconds.
125    pub(super) updated_at_ms: u128,
126    /// Optional progress counters/message.
127    pub(super) progress: Option<McpTaskProgress>,
128    /// Concise failure diagnostic when present.
129    pub(super) error: Option<String>,
130    /// Result reference or follow-up tool when present.
131    pub(super) result_ref: Option<String>,
132    /// Whether this task can be canceled by the current server.
133    pub(super) cancelable: bool,
134    /// Shared cooperative cancellation boundary for active indexing work.
135    #[serde(skip)]
136    pub(super) control: Option<IndexWorkControl>,
137}
138
139impl McpTaskRecord {
140    /// Return whether this record is in a terminal state and can be evicted first.
141    pub(super) fn is_terminal_state(&self) -> bool {
142        matches!(
143            self.state,
144            McpTaskState::Complete | McpTaskState::Failed | McpTaskState::Canceled
145        )
146    }
147}
148
149/// MCP task operation kind.
150#[derive(Debug, Clone, Eq, PartialEq, Serialize)]
151#[serde(rename_all = "snake_case")]
152pub(super) enum McpTaskOperation {
153    /// Contract/schema marker task.
154    Contract,
155    /// Repository scan and index operation.
156    Scan,
157    /// One-shot watch refresh operation.
158    WatchOnce,
159    /// Symbol projection rebuild operation.
160    SymbolsBuild,
161    /// Future search operation.
162    Search,
163}
164
165/// MCP task lifecycle state.
166#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize)]
167#[serde(rename_all = "snake_case")]
168pub(super) enum McpTaskState {
169    /// Task has not started.
170    Pending,
171    /// Task is running.
172    Running,
173    /// Task completed successfully.
174    Complete,
175    /// Task failed.
176    Failed,
177    /// Task was canceled.
178    Canceled,
179}
180
181/// Optional task progress fields.
182#[derive(Debug, Clone, Serialize)]
183pub(super) struct McpTaskProgress {
184    /// Completed unit count when known.
185    pub(super) current: Option<u64>,
186    /// Total unit count when known.
187    pub(super) total: Option<u64>,
188    /// Concise progress message.
189    pub(super) message: Option<String>,
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    #[test]
197    fn evictions_prefer_old_terminal_records() {
198        let mut registry = McpTaskRegistry {
199            records: VecDeque::new(),
200        };
201        registry.insert(McpTaskRecord {
202            task_id: "running-0".to_string(),
203            operation: McpTaskOperation::Search,
204            state: McpTaskState::Running,
205            created_at_ms: 0,
206            updated_at_ms: 0,
207            progress: None,
208            error: None,
209            result_ref: None,
210            cancelable: true,
211            control: None,
212        });
213        for index in 1..MCP_TASK_REGISTRY_CAPACITY {
214            registry.insert(McpTaskRecord {
215                task_id: format!("complete-{index}"),
216                operation: McpTaskOperation::Search,
217                state: McpTaskState::Complete,
218                created_at_ms: index as u128,
219                updated_at_ms: index as u128,
220                progress: None,
221                error: None,
222                result_ref: None,
223                cancelable: false,
224                control: None,
225            });
226        }
227
228        registry.insert(McpTaskRecord {
229            task_id: "new-complete".to_string(),
230            operation: McpTaskOperation::Search,
231            state: McpTaskState::Complete,
232            created_at_ms: 100,
233            updated_at_ms: 100,
234            progress: None,
235            error: None,
236            result_ref: Some(MCP_TOOL_ATLAS_TASK_STATUS.to_string()),
237            cancelable: false,
238            control: None,
239        });
240
241        assert_eq!(registry.len(), MCP_TASK_REGISTRY_CAPACITY);
242        assert!(registry.get("running-0").is_some());
243        assert!(registry.get("complete-1").is_none());
244        assert!(registry.get("new-complete").is_some());
245    }
246}
projectatlas/mcp/task_registry.rs

projectatlas/mcp/
task_registry.rs
