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 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_id: String,
44    ///         _request: CreateMessageRequest,
45    ///     ) -> Result<CreateMessageResult, Box<dyn std::error::Error + Send + Sync>> {
46    ///         // Handle sampling request (use request_id for tracking/correlation)
47    ///         todo!("Implement sampling logic")
48    ///     }
49    /// }
50    ///
51    /// let mut client = Client::new(StdioTransport::new());
52    /// client.set_sampling_handler(Arc::new(ExampleHandler));
53    /// ```
54    pub fn set_sampling_handler(&self, handler: Arc<dyn SamplingHandler>) {
55        *self
56            .inner
57            .sampling_handler
58            .lock()
59            .expect("sampling_handler mutex poisoned") = Some(handler);
60    }
61
62    /// Check if sampling is enabled
63    ///
64    /// Returns true if a sampling handler has been configured and sampling
65    /// capabilities are enabled.
66    pub fn has_sampling_handler(&self) -> bool {
67        self.inner
68            .sampling_handler
69            .lock()
70            .expect("sampling_handler mutex poisoned")
71            .is_some()
72    }
73
74    /// Remove the sampling handler
75    ///
76    /// Disables sampling capabilities and removes the handler. The client
77    /// will no longer advertise sampling support to servers.
78    pub fn remove_sampling_handler(&self) {
79        *self
80            .inner
81            .sampling_handler
82            .lock()
83            .expect("sampling_handler mutex poisoned") = None;
84    }
85
86    /// Get sampling capabilities for initialization
87    ///
88    /// Returns the sampling capabilities to be sent during client initialization
89    /// if sampling is enabled.
90    pub(crate) fn get_sampling_capabilities(&self) -> Option<SamplingCapabilities> {
91        if self
92            .inner
93            .sampling_handler
94            .lock()
95            .expect("sampling_handler mutex poisoned")
96            .is_some()
97        {
98            Some(SamplingCapabilities)
99        } else {
100            None
101        }
102    }
103}
turbomcp_client/client/operations/sampling.rs

turbomcp_client/client/operations/
sampling.rs
