mcpkit_server/capability/
tasks.rs

1//! Task capability implementation.
2//!
3//! Tasks let a server run a long-running operation while the caller polls for
4//! status (`tasks/get`) and, once terminal, the payload (`tasks/result`).
5//!
6//! The store itself ([`TaskManager`], [`TaskHandle`], [`route_task_store`]) is
7//! shared with the client side and lives in [`mcpkit_core::tasks`]; this
8//! module re-exports it and adds the server-only [`TaskService`].
9
10pub use mcpkit_core::tasks::{
11    DEFAULT_TASK_TTL_MS, RELATED_TASK_META_KEY, TaskEvent, TaskHandle, TaskManager, TaskObserver,
12    TaskPayload, TaskRoute, TaskState, route_task_store,
13};
14
15/// Where an ambient notification goes on a given transport.
16///
17/// A task transition has no request-scoped [`Peer`](crate::context::Peer) to
18/// send on, and each transport reaches its client differently: the stdio/socket
19/// runtime queues onto [`ServerState`](crate::server::ServerState)'s ambient
20/// pump, while the HTTP adapters store-and-forward onto the session's SSE
21/// [`StreamRegistry`](crate::streams::StreamRegistry).
22///
23/// The trait exists so the *mapping* from a task transition to
24/// `notifications/tasks/status` is written once. A new transport implements one
25/// method; it does not re-derive the notification, and so cannot get it subtly
26/// wrong the way five hand-rolled copies of a routing rule did.
27///
28/// Publishing is best-effort by contract: a notification with nowhere to go is
29/// dropped, never an error.
30pub trait NotificationSink: Send + Sync {
31    /// Hand the notification to this transport's outbound path.
32    fn publish(&self, notification: Notification);
33}
34
35impl NotificationSink for crate::server::ServerState {
36    fn publish(&self, notification: Notification) {
37        self.publish_notification(notification);
38    }
39}
40
41impl NotificationSink for crate::streams::StreamRegistry {
42    fn publish(&self, notification: Notification) {
43        match serde_json::to_string(&mcpkit_core::protocol::Message::Notification(notification)) {
44            // `None` simply means no live stream; the event is buffered for a
45            // resuming GET, and a client that never returns misses it.
46            Ok(json) => {
47                let _ = self.send("message", json);
48            }
49            Err(e) => tracing::warn!(error = ?e, "failed to serialize ambient notification"),
50        }
51    }
52}
53
54/// Publishes `notifications/tasks/status` when a task changes status.
55///
56/// The store emits domain events ([`TaskEvent`]); this is the only place that
57/// decides a transition is worth telling the client about, and the only place
58/// that builds the notification. Where it goes is the [`NotificationSink`]'s
59/// business.
60///
61/// Per spec the notification carries the task state in `params` and must **not**
62/// be tagged with `io.modelcontextprotocol/related-task` — the `taskId` is
63/// already there. [`TaskStatusNotificationParams`] carries no `_meta` when built
64/// from a [`Task`](mcpkit_core::types::task::Task), which is what keeps that true.
65pub struct TaskStatusNotifier {
66    sink: Arc<dyn NotificationSink>,
67}
68
69impl std::fmt::Debug for TaskStatusNotifier {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        // The sink is not `Debug` (implementors hold locks and capability
72        // sets); the observer's identity is all a reader needs here.
73        f.debug_struct("TaskStatusNotifier").finish_non_exhaustive()
74    }
75}
76
77impl TaskStatusNotifier {
78    /// Publish status transitions onto `sink`.
79    #[must_use]
80    pub const fn new(sink: Arc<dyn NotificationSink>) -> Self {
81        Self { sink }
82    }
83}
84
85impl TaskObserver for TaskStatusNotifier {
86    fn on_task_event(&self, event: &TaskEvent) {
87        let params = TaskStatusNotificationParams::from(event.task.clone());
88        match serde_json::to_value(params) {
89            Ok(params) => self.sink.publish(Notification::with_params(
90                crate::router::notifications::TASK_STATUS,
91                params,
92            )),
93            Err(e) => {
94                tracing::warn!(error = ?e, "failed to serialize task status notification");
95            }
96        }
97    }
98}
99
100/// Build a per-session task store that publishes `notifications/tasks/status`
101/// onto the session's SSE stream registry.
102///
103/// Every HTTP adapter creates its `TaskManager` and `StreamRegistry` together
104/// per session; this is the one place that wires them, so an adapter cannot
105/// forget to and silently stop emitting the notification.
106///
107/// `default_ttl_ms` mirrors [`TaskManager::with_default_ttl`].
108#[must_use]
109pub fn session_task_store(
110    streams: &Arc<crate::streams::StreamRegistry>,
111    default_ttl_ms: Option<u64>,
112) -> Arc<TaskManager> {
113    let store = Arc::new(TaskManager::with_default_ttl(default_ttl_ms));
114    // Only fails if an observer were already registered, which cannot happen on
115    // a store constructed a line ago.
116    let sink: Arc<dyn NotificationSink> = Arc::<crate::streams::StreamRegistry>::clone(streams);
117    let _ = store.set_observer(Arc::new(TaskStatusNotifier::new(sink)));
118    store
119}
120
121use crate::context::Context;
122use crate::handler::TaskHandler;
123use mcpkit_core::error::McpError;
124use mcpkit_core::protocol::Notification;
125use mcpkit_core::types::task::{
126    CancelTaskResult, GetTaskResult, ListTasksResult, TaskId, TaskStatusNotificationParams,
127};
128use std::sync::Arc;
129
130/// Task service implementing the [`TaskHandler`] trait over a [`TaskManager`].
131pub struct TaskService {
132    manager: Arc<TaskManager>,
133}
134
135impl Default for TaskService {
136    fn default() -> Self {
137        Self::new()
138    }
139}
140
141impl TaskService {
142    /// Create a new task service.
143    #[must_use]
144    pub fn new() -> Self {
145        Self {
146            manager: Arc::new(TaskManager::new()),
147        }
148    }
149
150    /// Get the underlying task manager.
151    #[must_use]
152    pub const fn manager(&self) -> &Arc<TaskManager> {
153        &self.manager
154    }
155
156    /// Create a new task and return a handle for driving it.
157    #[must_use]
158    pub fn create(&self) -> TaskHandle {
159        self.manager.create(None)
160    }
161}
162
163impl TaskHandler for TaskService {
164    async fn list_tasks(&self, _ctx: &Context<'_>) -> Result<ListTasksResult, McpError> {
165        Ok(self.manager.list().into())
166    }
167
168    async fn get_task(
169        &self,
170        task_id: &TaskId,
171        _ctx: &Context<'_>,
172    ) -> Result<Option<GetTaskResult>, McpError> {
173        Ok(self
174            .manager
175            .get(task_id)
176            .map(|s| GetTaskResult::from(s.task)))
177    }
178
179    async fn cancel_task(
180        &self,
181        task_id: &TaskId,
182        _ctx: &Context<'_>,
183    ) -> Result<Option<CancelTaskResult>, McpError> {
184        // Unknown task -> Ok(None); a real internal failure (e.g. poisoned lock)
185        // must surface as Err, not be collapsed into "unknown".
186        if self.manager.get(task_id).is_none() {
187            return Ok(None);
188        }
189        self.manager.cancel(task_id)?;
190        Ok(self
191            .manager
192            .get(task_id)
193            .map(|s| CancelTaskResult::from(s.task)))
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[tokio::test]
202    async fn test_task_service_handler() -> Result<(), Box<dyn std::error::Error>> {
203        let service = TaskService::new();
204        let handle = service.create();
205        let task_id = handle.id().clone();
206
207        assert_eq!(service.manager().list().len(), 1);
208        assert!(service.manager().get(&task_id).is_some());
209        Ok(())
210    }
211}
212
213#[cfg(test)]
214mod notifier_tests {
215    use super::*;
216    use crate::streams::{StreamConfig, StreamRegistry};
217
218    /// `session_task_store` is the one place adapters wire a store to a stream;
219    /// if it stops publishing, every HTTP transport silently stops emitting
220    /// `notifications/tasks/status`.
221    #[tokio::test]
222    async fn session_task_store_publishes_transitions_onto_the_registry() {
223        let streams = Arc::new(StreamRegistry::new(StreamConfig::default()));
224        let (mut handle, _prime) = streams.open("message", "{}".to_string());
225        let store = session_task_store(&streams, None);
226
227        let task = store.create(None);
228        task.complete(serde_json::json!({"ok": true}))
229            .expect("complete");
230
231        let event = handle.recv().await.expect("an event");
232        let json: serde_json::Value = serde_json::from_str(&event.data).expect("json");
233        assert_eq!(json["method"], "notifications/tasks/status");
234        assert_eq!(json["params"]["status"], "completed");
235        assert_eq!(json["params"]["taskId"], task.id().as_str());
236        assert!(
237            json["params"]["_meta"].is_null(),
238            "must not carry related-task _meta: {json}"
239        );
240    }
241
242    /// A store with no live stream must not error or panic — notifications are
243    /// best-effort by contract.
244    #[tokio::test]
245    async fn publishing_with_no_live_stream_is_a_no_op() {
246        let streams = Arc::new(StreamRegistry::new(StreamConfig::default()));
247        let store = session_task_store(&streams, None);
248        let task = store.create(None);
249        task.complete(serde_json::json!({})).expect("complete");
250        assert!(!streams.has_live_stream());
251    }
252}
mcpkit_server/capability/tasks.rs

mcpkit_server/capability/
tasks.rs
