turbomcp_client/client/operations/
sampling.rs

1//! Sampling operations for MCP client
2//!
3//! This module provides sampling capability management for LLM operations.
4//! Sampling allows the MCP server to request the client to perform LLM
5//! inference when the server needs language model capabilities.
6//!
7//! The client's role in sampling is to:
8//! 1. Register handlers for sampling/createMessage requests
9//! 2. Advertise sampling capabilities during initialization
10//! 3. Process server-initiated sampling requests (handled in core message routing)
11
12use crate::sampling::SamplingHandler;
13use std::sync::Arc;
14use turbomcp_protocol::types::SamplingCapabilities;
15
16impl<T: turbomcp_transport::Transport> super::super::core::Client<T> {
17    /// Set the sampling handler for processing server-initiated sampling requests
18    ///
19    /// Registers a handler that can process LLM sampling requests from the server.
20    /// When a handler is set, the client will advertise sampling capabilities
21    /// during initialization, allowing the server to request LLM operations.
22    ///
23    /// # Arguments
24    ///
25    /// * `handler` - The handler implementation for sampling requests
26    ///
27    /// # Examples
28    ///
29    /// ```rust,no_run
30    /// use turbomcp_client::{Client, sampling::SamplingHandler};
31    /// use turbomcp_transport::stdio::StdioTransport;
32    /// use turbomcp_protocol::types::{CreateMessageRequest, CreateMessageResult};
33    /// use async_trait::async_trait;
34    /// use std::sync::Arc;
35    ///
36    /// #[derive(Debug)]
37    /// struct ExampleHandler;
38    ///
39    /// #[async_trait]
40    /// impl SamplingHandler for ExampleHandler {
41    ///     async fn handle_create_message(
42    ///         &self,
43    ///         _request: CreateMessageRequest,
44    ///     ) -> Result<CreateMessageResult, Box<dyn std::error::Error + Send + Sync>> {
45    ///         // Handle sampling request
46    ///         todo!("Implement sampling logic")
47    ///     }
48    /// }
49    ///
50    /// let mut client = Client::new(StdioTransport::new());
51    /// client.set_sampling_handler(Arc::new(ExampleHandler));
52    /// ```
53    pub fn set_sampling_handler(&self, handler: Arc<dyn SamplingHandler>) {
54        *self
55            .inner
56            .sampling_handler
57            .lock()
58            .expect("sampling_handler mutex poisoned") = Some(handler);
59    }
60
61    /// Check if sampling is enabled
62    ///
63    /// Returns true if a sampling handler has been configured and sampling
64    /// capabilities are enabled.
65    pub fn has_sampling_handler(&self) -> bool {
66        self.inner
67            .sampling_handler
68            .lock()
69            .expect("sampling_handler mutex poisoned")
70            .is_some()
71    }
72
73    /// Remove the sampling handler
74    ///
75    /// Disables sampling capabilities and removes the handler. The client
76    /// will no longer advertise sampling support to servers.
77    pub fn remove_sampling_handler(&self) {
78        *self
79            .inner
80            .sampling_handler
81            .lock()
82            .expect("sampling_handler mutex poisoned") = None;
83    }
84
85    /// Get sampling capabilities for initialization
86    ///
87    /// Returns the sampling capabilities to be sent during client initialization
88    /// if sampling is enabled.
89    pub(crate) fn get_sampling_capabilities(&self) -> Option<SamplingCapabilities> {
90        if self
91            .inner
92            .sampling_handler
93            .lock()
94            .expect("sampling_handler mutex poisoned")
95            .is_some()
96        {
97            Some(SamplingCapabilities)
98        } else {
99            None
100        }
101    }
102}
turbomcp_client/client/operations/sampling.rs

turbomcp_client/client/operations/
sampling.rs
