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}