turbomcp_server/
sampling.rs1use crate::{ServerError, ServerResult};
7use turbomcp_protocol::RequestContext;
8use turbomcp_protocol::types::{
9 CreateMessageRequest, CreateMessageResult, ElicitRequest, ElicitResult, ListRootsResult,
10};
11
12#[async_trait::async_trait]
16pub trait SamplingExt {
17 async fn create_message(
19 &self,
20 request: CreateMessageRequest,
21 ) -> ServerResult<CreateMessageResult>;
22
23 async fn elicit(&self, request: ElicitRequest) -> ServerResult<ElicitResult>;
25
26 async fn list_roots(&self) -> ServerResult<ListRootsResult>;
28}
29
30#[async_trait::async_trait]
31impl SamplingExt for RequestContext {
32 async fn create_message(
33 &self,
34 request: CreateMessageRequest,
35 ) -> ServerResult<CreateMessageResult> {
36 let capabilities = self
37 .server_to_client()
38 .ok_or_else(|| ServerError::Handler {
39 message: "No server capabilities available for sampling requests".to_string(),
40 context: Some("sampling".to_string()),
41 })?;
42
43 capabilities
45 .create_message(request, self.clone())
46 .await
47 .map_err(|e| ServerError::Handler {
48 message: format!("Sampling request failed: {}", e),
49 context: Some("sampling".to_string()),
50 })
51 }
52
53 async fn elicit(&self, request: ElicitRequest) -> ServerResult<ElicitResult> {
54 let capabilities = self
55 .server_to_client()
56 .ok_or_else(|| ServerError::Handler {
57 message: "No server capabilities available for elicitation requests".to_string(),
58 context: Some("elicitation".to_string()),
59 })?;
60
61 capabilities
63 .elicit(request, self.clone())
64 .await
65 .map_err(|e| ServerError::Handler {
66 message: format!("Elicitation request failed: {}", e),
67 context: Some("elicitation".to_string()),
68 })
69 }
70
71 async fn list_roots(&self) -> ServerResult<ListRootsResult> {
72 let capabilities = self
73 .server_to_client()
74 .ok_or_else(|| ServerError::Handler {
75 message: "No server capabilities available for roots listing".to_string(),
76 context: Some("roots".to_string()),
77 })?;
78
79 capabilities
81 .list_roots(self.clone())
82 .await
83 .map_err(|e| ServerError::Handler {
84 message: format!("Roots listing failed: {}", e),
85 context: Some("roots".to_string()),
86 })
87 }
88}