Skip to main content

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 + 'static> 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 std::sync::Arc;
34    /// use std::future::Future;
35    /// use std::pin::Pin;
36    ///
37    /// #[derive(Debug)]
38    /// struct ExampleHandler;
39    ///
40    /// impl SamplingHandler for ExampleHandler {
41    ///     fn handle_create_message(
42    ///         &self,
43    ///         _request_id: String,
44    ///         _request: CreateMessageRequest,
45    ///     ) -> Pin<Box<dyn Future<Output = Result<CreateMessageResult, Box<dyn std::error::Error + Send + Sync>>> + Send + '_>> {
46    ///         Box::pin(async move {
47    ///             // Handle sampling request (use request_id for tracking/correlation)
48    ///             todo!("Implement sampling logic")
49    ///         })
50    ///     }
51    /// }
52    ///
53    /// let mut client = Client::new(StdioTransport::new());
54    /// client.set_sampling_handler(Arc::new(ExampleHandler));
55    /// ```
56    pub fn set_sampling_handler(&self, handler: Arc<dyn SamplingHandler>) {
57        *self.inner.sampling_handler.lock() = Some(handler);
58    }
59
60    /// Check if sampling is enabled
61    ///
62    /// Returns true if a sampling handler has been configured and sampling
63    /// capabilities are enabled.
64    #[must_use]
65    pub fn has_sampling_handler(&self) -> bool {
66        self.inner.sampling_handler.lock().is_some()
67    }
68
69    /// Remove the sampling handler
70    ///
71    /// Disables sampling capabilities and removes the handler. The client
72    /// will no longer advertise sampling support to servers.
73    pub fn remove_sampling_handler(&self) {
74        *self.inner.sampling_handler.lock() = None;
75    }
76
77    /// Get sampling capabilities for initialization
78    ///
79    /// Returns the sampling capabilities to be sent during client initialization
80    /// if sampling is enabled.
81    pub(crate) fn get_sampling_capabilities(&self) -> Option<SamplingCapabilities> {
82        if self.inner.sampling_handler.lock().is_some() {
83            Some(SamplingCapabilities {})
84        } else {
85            None
86        }
87    }
88}