tower_mcp/client/
handler.rs

1//! Handler trait for server-initiated requests and notifications.
2//!
3//! The [`ClientHandler`] trait defines how the client responds to requests
4//! and notifications sent by the server. All methods have default
5//! implementations, so you only need to override the ones you care about.
6//!
7//! The unit type `()` implements this trait with all defaults, which is
8//! used by [`McpClient::connect()`](super::McpClient::connect).
9//!
10//! # Notification Handler
11//!
12//! For notification-only use cases, [`NotificationHandler`] provides a
13//! builder-based alternative to implementing the full trait:
14//!
15//! ```rust
16//! use tower_mcp::client::NotificationHandler;
17//!
18//! let handler = NotificationHandler::new()
19//!     .on_tools_changed(|| {
20//!         println!("Tools changed, re-fetching...");
21//!     })
22//!     .on_log_message(|msg| {
23//!         println!("[{}] {}", msg.level, msg.data);
24//!     });
25//! ```
26//!
27//! For forwarding MCP log messages to the [`tracing`] crate:
28//!
29//! ```rust
30//! use tower_mcp::client::NotificationHandler;
31//!
32//! let handler = NotificationHandler::with_log_forwarding();
33//! ```
34//!
35//! # Custom Handler
36//!
37//! ```rust,ignore
38//! use async_trait::async_trait;
39//! use tower_mcp::client::ClientHandler;
40//! use tower_mcp::protocol::{CreateMessageParams, CreateMessageResult};
41//! use tower_mcp_types::JsonRpcError;
42//!
43//! struct MySamplingHandler;
44//!
45//! #[async_trait]
46//! impl ClientHandler for MySamplingHandler {
47//!     async fn handle_create_message(
48//!         &self,
49//!         params: CreateMessageParams,
50//!     ) -> Result<CreateMessageResult, JsonRpcError> {
51//!         // Forward to your LLM and return the result
52//!         todo!()
53//!     }
54//! }
55//! ```
56
57use async_trait::async_trait;
58
59use crate::protocol::{
60    CreateMessageParams, CreateMessageResult, ElicitRequestParams, ElicitResult, ListRootsResult,
61    LogLevel, LoggingMessageParams, ProgressParams, RequestId, SubscriptionFilter,
62    TaskStatusParams,
63};
64use crate::tasks::TaskStatusNotificationParams;
65use tower_mcp_types::JsonRpcError;
66
67/// Notification sent from the server to the client.
68///
69/// These correspond to the `notifications/` methods defined in the MCP spec
70/// that flow from server to client.
71#[derive(Debug, Clone)]
72#[non_exhaustive]
73pub enum ServerNotification {
74    /// Progress update for a request (`notifications/progress`).
75    Progress(ProgressParams),
76    /// Log message (`notifications/message`).
77    LogMessage(LoggingMessageParams),
78    /// A subscribed resource has been updated (`notifications/resources/updated`).
79    ResourceUpdated {
80        /// The URI of the updated resource.
81        uri: String,
82    },
83    /// The list of available resources has changed.
84    ResourcesListChanged,
85    /// The list of available tools has changed.
86    ToolsListChanged,
87    /// The list of available prompts has changed.
88    PromptsListChanged,
89    /// A legacy task changed status (`notifications/tasks`).
90    TaskStatusChanged(TaskStatusParams),
91    /// A final-protocol task changed status (`notifications/tasks`).
92    FinalTaskStatusChanged(TaskStatusNotificationParams),
93    /// The server acknowledged a `subscriptions/listen` stream.
94    SubscriptionAcknowledged {
95        /// JSON-RPC request ID that identifies the subscription.
96        subscription_id: RequestId,
97        /// Subset of the requested filter the server agreed to honor.
98        notifications: SubscriptionFilter,
99    },
100    /// A notification delivered on a `subscriptions/listen` stream.
101    Subscription {
102        /// JSON-RPC request ID that identifies the subscription.
103        subscription_id: RequestId,
104        /// The ordinary notification carried by this subscription.
105        notification: Box<ServerNotification>,
106    },
107    /// The server cancelled an active subscription.
108    SubscriptionCancelled {
109        /// JSON-RPC request ID that identified the subscription.
110        subscription_id: RequestId,
111        /// Optional diagnostic reason from the server.
112        reason: Option<String>,
113    },
114    /// An unknown or unrecognized notification.
115    Unknown {
116        /// The notification method name.
117        method: String,
118        /// The notification parameters, if any.
119        params: Option<serde_json::Value>,
120    },
121}
122
123/// Handler for server-initiated requests and notifications.
124///
125/// Implement this trait to handle sampling requests, elicitation requests,
126/// roots listing, and server notifications. All methods have default
127/// implementations that either return sensible defaults or reject with
128/// `method_not_found`.
129///
130/// The unit type `()` implements this trait with all defaults, which is
131/// used by [`McpClient::connect()`](super::McpClient::connect).
132#[async_trait]
133pub trait ClientHandler: Send + Sync + 'static {
134    /// Handle a `sampling/createMessage` request from the server.
135    ///
136    /// The server is asking the client to perform LLM inference. Override
137    /// this to forward the request to your LLM provider.
138    ///
139    /// Default: returns `method_not_found` error.
140    async fn handle_create_message(
141        &self,
142        _params: CreateMessageParams,
143    ) -> Result<CreateMessageResult, JsonRpcError> {
144        Err(JsonRpcError::method_not_found("sampling/createMessage"))
145    }
146
147    /// Handle an `elicitation/create` request from the server.
148    ///
149    /// The server is asking the client for user input (form data or URL).
150    ///
151    /// Default: returns `method_not_found` error.
152    async fn handle_elicit(
153        &self,
154        _params: ElicitRequestParams,
155    ) -> Result<ElicitResult, JsonRpcError> {
156        Err(JsonRpcError::method_not_found("elicitation/create"))
157    }
158
159    /// Handle a `roots/list` request from the server.
160    ///
161    /// The server is asking which filesystem roots the client has access to.
162    /// If roots were configured on the [`McpClient`](super::McpClient) via
163    /// the builder, those are returned automatically before this method
164    /// is called.
165    ///
166    /// Default: returns an empty list.
167    async fn handle_list_roots(&self) -> Result<ListRootsResult, JsonRpcError> {
168        Ok(ListRootsResult {
169            roots: vec![],
170            meta: None,
171        })
172    }
173
174    /// Called when the server sends a notification.
175    ///
176    /// Override to handle progress updates, log messages, resource changes, etc.
177    ///
178    /// Default: no-op.
179    async fn on_notification(&self, _notification: ServerNotification) {}
180}
181
182/// Unit type implements [`ClientHandler`] with all defaults.
183#[async_trait]
184impl ClientHandler for () {}
185
186// Type aliases for notification callback boxes.
187type ProgressCallback = Box<dyn Fn(ProgressParams) + Send + Sync>;
188type LogMessageCallback = Box<dyn Fn(LoggingMessageParams) + Send + Sync>;
189type ResourceUpdatedCallback = Box<dyn Fn(String) + Send + Sync>;
190type TaskStatusCallback = Box<dyn Fn(TaskStatusParams) + Send + Sync>;
191type FinalTaskStatusCallback = Box<dyn Fn(TaskStatusNotificationParams) + Send + Sync>;
192type SimpleCallback = Box<dyn Fn() + Send + Sync>;
193
194/// Callback-based handler for server notifications.
195///
196/// Provides typed callback registration for each notification type,
197/// without requiring a full [`ClientHandler`] trait implementation.
198/// Server-initiated requests (sampling, elicitation, roots) are
199/// rejected with `method_not_found`.
200///
201/// # Example
202///
203/// ```rust
204/// use tower_mcp::client::NotificationHandler;
205///
206/// let handler = NotificationHandler::new()
207///     .on_progress(|p| {
208///         println!("Progress: {}/{}", p.progress, p.total.unwrap_or(1.0));
209///     })
210///     .on_tools_changed(|| {
211///         println!("Server tools changed!");
212///     });
213/// ```
214pub struct NotificationHandler {
215    on_progress: Option<ProgressCallback>,
216    on_log_message: Option<LogMessageCallback>,
217    on_resource_updated: Option<ResourceUpdatedCallback>,
218    on_resources_changed: Option<SimpleCallback>,
219    on_tools_changed: Option<SimpleCallback>,
220    on_prompts_changed: Option<SimpleCallback>,
221    on_task_status_changed: Option<TaskStatusCallback>,
222    on_final_task_status_changed: Option<FinalTaskStatusCallback>,
223}
224
225impl NotificationHandler {
226    /// Create a new handler with no callbacks registered.
227    pub fn new() -> Self {
228        Self {
229            on_progress: None,
230            on_log_message: None,
231            on_resource_updated: None,
232            on_resources_changed: None,
233            on_tools_changed: None,
234            on_prompts_changed: None,
235            on_task_status_changed: None,
236            on_final_task_status_changed: None,
237        }
238    }
239
240    /// Create a handler that forwards MCP log messages to [`tracing`].
241    ///
242    /// Maps MCP log levels to tracing levels:
243    /// - Emergency, Alert, Critical -> `error!`
244    /// - Error -> `error!`
245    /// - Warning -> `warn!`
246    /// - Notice, Info -> `info!`
247    /// - Debug -> `debug!`
248    pub fn with_log_forwarding() -> Self {
249        Self::new().on_log_message(|msg| {
250            let logger = msg.logger.as_deref().unwrap_or("mcp");
251            match msg.level {
252                LogLevel::Emergency | LogLevel::Alert | LogLevel::Critical | LogLevel::Error => {
253                    tracing::error!(logger = logger, "{}", msg.data);
254                }
255                LogLevel::Warning => {
256                    tracing::warn!(logger = logger, "{}", msg.data);
257                }
258                LogLevel::Notice | LogLevel::Info => {
259                    tracing::info!(logger = logger, "{}", msg.data);
260                }
261                LogLevel::Debug => {
262                    tracing::debug!(logger = logger, "{}", msg.data);
263                }
264                _ => {
265                    tracing::trace!(logger = logger, "{}", msg.data);
266                }
267            }
268        })
269    }
270
271    /// Register a callback for progress notifications.
272    pub fn on_progress(mut self, f: impl Fn(ProgressParams) + Send + Sync + 'static) -> Self {
273        self.on_progress = Some(Box::new(f));
274        self
275    }
276
277    /// Register a callback for log message notifications.
278    pub fn on_log_message(
279        mut self,
280        f: impl Fn(LoggingMessageParams) + Send + Sync + 'static,
281    ) -> Self {
282        self.on_log_message = Some(Box::new(f));
283        self
284    }
285
286    /// Register a callback for resource updated notifications.
287    ///
288    /// The callback receives the URI of the updated resource.
289    pub fn on_resource_updated(mut self, f: impl Fn(String) + Send + Sync + 'static) -> Self {
290        self.on_resource_updated = Some(Box::new(f));
291        self
292    }
293
294    /// Register a callback for resources list changed notifications.
295    pub fn on_resources_changed(mut self, f: impl Fn() + Send + Sync + 'static) -> Self {
296        self.on_resources_changed = Some(Box::new(f));
297        self
298    }
299
300    /// Register a callback for tools list changed notifications.
301    pub fn on_tools_changed(mut self, f: impl Fn() + Send + Sync + 'static) -> Self {
302        self.on_tools_changed = Some(Box::new(f));
303        self
304    }
305
306    /// Register a callback for prompts list changed notifications.
307    pub fn on_prompts_changed(mut self, f: impl Fn() + Send + Sync + 'static) -> Self {
308        self.on_prompts_changed = Some(Box::new(f));
309        self
310    }
311
312    /// Register a callback for legacy task status notifications.
313    pub fn on_task_status_changed(
314        mut self,
315        f: impl Fn(TaskStatusParams) + Send + Sync + 'static,
316    ) -> Self {
317        self.on_task_status_changed = Some(Box::new(f));
318        self
319    }
320
321    /// Register a callback for final-protocol task status notifications.
322    pub fn on_final_task_status_changed(
323        mut self,
324        f: impl Fn(TaskStatusNotificationParams) + Send + Sync + 'static,
325    ) -> Self {
326        self.on_final_task_status_changed = Some(Box::new(f));
327        self
328    }
329
330    fn dispatch_notification(&self, notification: ServerNotification) {
331        match notification {
332            ServerNotification::Progress(params) => {
333                if let Some(cb) = &self.on_progress {
334                    cb(params);
335                }
336            }
337            ServerNotification::LogMessage(params) => {
338                if let Some(cb) = &self.on_log_message {
339                    cb(params);
340                }
341            }
342            ServerNotification::ResourceUpdated { uri } => {
343                if let Some(cb) = &self.on_resource_updated {
344                    cb(uri);
345                }
346            }
347            ServerNotification::ResourcesListChanged => {
348                if let Some(cb) = &self.on_resources_changed {
349                    cb();
350                }
351            }
352            ServerNotification::ToolsListChanged => {
353                if let Some(cb) = &self.on_tools_changed {
354                    cb();
355                }
356            }
357            ServerNotification::PromptsListChanged => {
358                if let Some(cb) = &self.on_prompts_changed {
359                    cb();
360                }
361            }
362            ServerNotification::TaskStatusChanged(params) => {
363                if let Some(cb) = &self.on_task_status_changed {
364                    cb(params);
365                }
366            }
367            ServerNotification::FinalTaskStatusChanged(params) => {
368                if let Some(cb) = &self.on_final_task_status_changed {
369                    cb(params);
370                }
371            }
372            // Preserve the existing callback API for subscription-delivered
373            // events while custom ClientHandler implementations can inspect
374            // the wrapper and correlate concurrent streams.
375            ServerNotification::Subscription { notification, .. } => {
376                self.dispatch_notification(*notification);
377            }
378            ServerNotification::SubscriptionAcknowledged { .. }
379            | ServerNotification::SubscriptionCancelled { .. }
380            | ServerNotification::Unknown { .. } => {}
381        }
382    }
383}
384
385impl Default for NotificationHandler {
386    fn default() -> Self {
387        Self::new()
388    }
389}
390
391impl std::fmt::Debug for NotificationHandler {
392    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
393        f.debug_struct("NotificationHandler")
394            .field("on_progress", &self.on_progress.is_some())
395            .field("on_log_message", &self.on_log_message.is_some())
396            .field("on_resource_updated", &self.on_resource_updated.is_some())
397            .field("on_resources_changed", &self.on_resources_changed.is_some())
398            .field("on_tools_changed", &self.on_tools_changed.is_some())
399            .field("on_prompts_changed", &self.on_prompts_changed.is_some())
400            .field(
401                "on_task_status_changed",
402                &self.on_task_status_changed.is_some(),
403            )
404            .field(
405                "on_final_task_status_changed",
406                &self.on_final_task_status_changed.is_some(),
407            )
408            .finish()
409    }
410}
411
412#[async_trait]
413impl ClientHandler for NotificationHandler {
414    async fn on_notification(&self, notification: ServerNotification) {
415        self.dispatch_notification(notification);
416    }
417}
418
419#[cfg(test)]
420mod tests {
421    use super::*;
422    use std::sync::Arc;
423    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
424
425    #[tokio::test]
426    async fn test_notification_handler_progress() {
427        let called = Arc::new(AtomicBool::new(false));
428        let called_clone = called.clone();
429        let handler = NotificationHandler::new().on_progress(move |p| {
430            assert!((p.progress - 0.5).abs() < f64::EPSILON);
431            called_clone.store(true, Ordering::SeqCst);
432        });
433
434        handler
435            .on_notification(ServerNotification::Progress(ProgressParams {
436                progress_token: crate::protocol::ProgressToken::String("t1".into()),
437                progress: 0.5,
438                total: Some(1.0),
439                message: None,
440                meta: None,
441            }))
442            .await;
443
444        assert!(called.load(Ordering::SeqCst));
445    }
446
447    #[tokio::test]
448    async fn test_notification_handler_log_message() {
449        let called = Arc::new(AtomicBool::new(false));
450        let called_clone = called.clone();
451        let handler = NotificationHandler::new().on_log_message(move |msg| {
452            assert_eq!(msg.level, LogLevel::Info);
453            called_clone.store(true, Ordering::SeqCst);
454        });
455
456        handler
457            .on_notification(ServerNotification::LogMessage(LoggingMessageParams {
458                level: LogLevel::Info,
459                logger: Some("test".into()),
460                data: serde_json::json!("hello"),
461                meta: None,
462            }))
463            .await;
464
465        assert!(called.load(Ordering::SeqCst));
466    }
467
468    #[tokio::test]
469    async fn test_notification_handler_resource_updated() {
470        let called = Arc::new(AtomicBool::new(false));
471        let called_clone = called.clone();
472        let handler = NotificationHandler::new().on_resource_updated(move |uri| {
473            assert_eq!(uri, "file:///test.txt");
474            called_clone.store(true, Ordering::SeqCst);
475        });
476
477        handler
478            .on_notification(ServerNotification::ResourceUpdated {
479                uri: "file:///test.txt".to_string(),
480            })
481            .await;
482
483        assert!(called.load(Ordering::SeqCst));
484    }
485
486    #[tokio::test]
487    async fn test_notification_handler_list_changed() {
488        let tools_count = Arc::new(AtomicUsize::new(0));
489        let resources_count = Arc::new(AtomicUsize::new(0));
490        let prompts_count = Arc::new(AtomicUsize::new(0));
491
492        let tc = tools_count.clone();
493        let rc = resources_count.clone();
494        let pc = prompts_count.clone();
495
496        let handler = NotificationHandler::new()
497            .on_tools_changed(move || {
498                tc.fetch_add(1, Ordering::SeqCst);
499            })
500            .on_resources_changed(move || {
501                rc.fetch_add(1, Ordering::SeqCst);
502            })
503            .on_prompts_changed(move || {
504                pc.fetch_add(1, Ordering::SeqCst);
505            });
506
507        handler
508            .on_notification(ServerNotification::ToolsListChanged)
509            .await;
510        handler
511            .on_notification(ServerNotification::ResourcesListChanged)
512            .await;
513        handler
514            .on_notification(ServerNotification::PromptsListChanged)
515            .await;
516
517        assert_eq!(tools_count.load(Ordering::SeqCst), 1);
518        assert_eq!(resources_count.load(Ordering::SeqCst), 1);
519        assert_eq!(prompts_count.load(Ordering::SeqCst), 1);
520    }
521
522    #[tokio::test]
523    async fn test_notification_handler_task_status_changed() {
524        let legacy_count = Arc::new(AtomicUsize::new(0));
525        let final_count = Arc::new(AtomicUsize::new(0));
526        let legacy = legacy_count.clone();
527        let final_ = final_count.clone();
528        let handler = NotificationHandler::new()
529            .on_task_status_changed(move |params| {
530                assert_eq!(params.task_id, "legacy-task");
531                legacy.fetch_add(1, Ordering::SeqCst);
532            })
533            .on_final_task_status_changed(move |params| {
534                assert_eq!(params.task.task_id(), "final-task");
535                final_.fetch_add(1, Ordering::SeqCst);
536            });
537
538        handler
539            .on_notification(ServerNotification::TaskStatusChanged(TaskStatusParams {
540                task_id: "legacy-task".into(),
541                status: crate::protocol::TaskStatus::Completed,
542                status_message: None,
543                created_at: "2026-08-02T00:00:00Z".into(),
544                last_updated_at: "2026-08-02T00:00:01Z".into(),
545                ttl: None,
546                poll_interval: None,
547                meta: None,
548            }))
549            .await;
550        handler
551            .on_notification(ServerNotification::FinalTaskStatusChanged(
552                TaskStatusNotificationParams {
553                    task: crate::tasks::DetailedTask::cancelled(crate::tasks::TaskMetadata::new(
554                        "final-task",
555                        "2026-08-02T00:00:00Z",
556                        "2026-08-02T00:00:01Z",
557                        None,
558                    )),
559                    meta: None,
560                },
561            ))
562            .await;
563
564        assert_eq!(legacy_count.load(Ordering::SeqCst), 1);
565        assert_eq!(final_count.load(Ordering::SeqCst), 1);
566    }
567
568    #[tokio::test]
569    async fn test_notification_handler_unset_callbacks_are_noop() {
570        // Handler with no callbacks should not panic
571        let handler = NotificationHandler::new();
572
573        handler
574            .on_notification(ServerNotification::ToolsListChanged)
575            .await;
576        handler
577            .on_notification(ServerNotification::Progress(ProgressParams {
578                progress_token: crate::protocol::ProgressToken::String("t".into()),
579                progress: 1.0,
580                total: None,
581                message: None,
582                meta: None,
583            }))
584            .await;
585        handler
586            .on_notification(ServerNotification::LogMessage(LoggingMessageParams {
587                level: LogLevel::Debug,
588                logger: None,
589                data: serde_json::json!("test"),
590                meta: None,
591            }))
592            .await;
593        handler
594            .on_notification(ServerNotification::Unknown {
595                method: "custom/thing".into(),
596                params: None,
597            })
598            .await;
599    }
600
601    #[tokio::test]
602    async fn test_notification_handler_rejects_requests() {
603        use crate::protocol::{ElicitFormParams, ElicitFormSchema};
604
605        let handler = NotificationHandler::new();
606
607        let params = serde_json::from_value::<CreateMessageParams>(serde_json::json!({
608            "messages": [],
609            "maxTokens": 100
610        }))
611        .unwrap();
612        let err = handler.handle_create_message(params).await.unwrap_err();
613        assert_eq!(err.code, -32601); // method_not_found
614
615        let err = handler
616            .handle_elicit(ElicitRequestParams::Form(ElicitFormParams {
617                mode: None,
618                message: "test".into(),
619                requested_schema: ElicitFormSchema {
620                    schema_type: "object".into(),
621                    properties: Default::default(),
622                    required: vec![],
623                },
624                meta: None,
625            }))
626            .await
627            .unwrap_err();
628        assert_eq!(err.code, -32601);
629    }
630
631    #[test]
632    fn test_notification_handler_debug() {
633        let handler = NotificationHandler::new().on_progress(|_| {});
634        let debug = format!("{:?}", handler);
635        assert!(debug.contains("on_progress: true"));
636        assert!(debug.contains("on_log_message: false"));
637    }
638
639    #[test]
640    fn test_notification_handler_default() {
641        let _handler = NotificationHandler::default();
642    }
643}
tower_mcp/client/handler.rs

tower_mcp/client/
handler.rs
