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    #[must_use]
67    pub fn has_sampling_handler(&self) -> bool {
68        self.inner
69            .sampling_handler
70            .lock()
71            .expect("sampling_handler mutex poisoned")
72            .is_some()
73    }
74
75    /// Remove the sampling handler
76    ///
77    /// Disables sampling capabilities and removes the handler. The client
78    /// will no longer advertise sampling support to servers.
79    pub fn remove_sampling_handler(&self) {
80        *self
81            .inner
82            .sampling_handler
83            .lock()
84            .expect("sampling_handler mutex poisoned") = None;
85    }
86
87    /// Get sampling capabilities for initialization
88    ///
89    /// Returns the sampling capabilities to be sent during client initialization
90    /// if sampling is enabled.
91    pub(crate) fn get_sampling_capabilities(&self) -> Option<SamplingCapabilities> {
92        if self
93            .inner
94            .sampling_handler
95            .lock()
96            .expect("sampling_handler mutex poisoned")
97            .is_some()
98        {
99            Some(SamplingCapabilities)
100        } else {
101            None
102        }
103    }
104}
turbomcp_client/client/operations/sampling.rs

turbomcp_client/client/operations/
sampling.rs
