feat: implement SEP-1577 sampling with tools support

Author: DaleSeo

crates/rmcp/src/model.rs

@@ -1209,18 +1209,257 @@ pub enum Role {
     Assistant,
 }
 
-/// A message in a sampling conversation, containing a role and content.
-///
-/// This represents a single message in a conversation flow, used primarily
-/// in LLM sampling requests where the conversation history is important
-/// for generating appropriate responses.
+/// Tool selection mode (SEP-1577).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub enum ToolChoiceMode {
+    /// Model decides whether to use tools
+    Auto,
+    /// Model must use at least one tool
+    Required,
+    /// Model must not use tools
+    None,
+}
+
+impl Default for ToolChoiceMode {
+    fn default() -> Self {
+        Self::Auto
+    }
+}
+
+/// Tool choice configuration (SEP-1577).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct ToolChoice {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub mode: Option<ToolChoiceMode>,
+}
+
+impl ToolChoice {
+    pub fn auto() -> Self {
+        Self {
+            mode: Some(ToolChoiceMode::Auto),
+        }
+    }
+
+    pub fn required() -> Self {
+        Self {
+            mode: Some(ToolChoiceMode::Required),
+        }
+    }
+
+    pub fn none() -> Self {
+        Self {
+            mode: Some(ToolChoiceMode::None),
+        }
+    }
+}
+
+/// Single or array content wrapper (SEP-1577).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub enum SamplingContent<T> {
+    Single(T),
+    Multiple(Vec<T>),
+}
+
+impl<T> SamplingContent<T> {
+    /// Convert to a Vec regardless of whether it's single or multiple
+    pub fn into_vec(self) -> Vec<T> {
+        match self {
+            SamplingContent::Single(item) => vec![item],
+            SamplingContent::Multiple(items) => items,
+        }
+    }
+
+    /// Check if the content is empty
+    pub fn is_empty(&self) -> bool {
+        match self {
+            SamplingContent::Single(_) => false,
+            SamplingContent::Multiple(items) => items.is_empty(),
+        }
+    }
+
+    /// Get the number of content items
+    pub fn len(&self) -> usize {
+        match self {
+            SamplingContent::Single(_) => 1,
+            SamplingContent::Multiple(items) => items.len(),
+        }
+    }
+}
+
+impl<T> Default for SamplingContent<T> {
+    fn default() -> Self {
+        SamplingContent::Multiple(Vec::new())
+    }
+}
+
+impl<T> SamplingContent<T> {
+    /// Get the first item if present
+    pub fn first(&self) -> Option<&T> {
+        match self {
+            SamplingContent::Single(item) => Some(item),
+            SamplingContent::Multiple(items) => items.first(),
+        }
+    }
+
+    /// Iterate over all content items
+    pub fn iter(&self) -> impl Iterator<Item = &T> {
+        let items: Vec<&T> = match self {
+            SamplingContent::Single(item) => vec![item],
+            SamplingContent::Multiple(items) => items.iter().collect(),
+        };
+        items.into_iter()
+    }
+}
+
+impl SamplingMessageContent {
+    /// Get the text content if this is a Text variant
+    pub fn as_text(&self) -> Option<&RawTextContent> {
+        match self {
+            SamplingMessageContent::Text(text) => Some(text),
+            _ => None,
+        }
+    }
+
+    /// Get the tool use content if this is a ToolUse variant
+    pub fn as_tool_use(&self) -> Option<&ToolUseContent> {
+        match self {
+            SamplingMessageContent::ToolUse(tool_use) => Some(tool_use),
+            _ => None,
+        }
+    }
+
+    /// Get the tool result content if this is a ToolResult variant
+    pub fn as_tool_result(&self) -> Option<&ToolResultContent> {
+        match self {
+            SamplingMessageContent::ToolResult(tool_result) => Some(tool_result),
+            _ => None,
+        }
+    }
+}
+
+impl<T> From<T> for SamplingContent<T> {
+    fn from(item: T) -> Self {
+        SamplingContent::Single(item)
+    }
+}
+
+impl<T> From<Vec<T>> for SamplingContent<T> {
+    fn from(items: Vec<T>) -> Self {
+        SamplingContent::Multiple(items)
+    }
+}
+
+/// A message in a sampling conversation (SEP-1577).
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
 pub struct SamplingMessage {
-    /// The role of the message sender (User or Assistant)
     pub role: Role,
-    /// The actual content of the message (text, image, etc.)
-    pub content: Content,
+    /// Single content or array of contents
+    pub content: SamplingContent<SamplingMessageContent>,
+    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
+    pub meta: Option<Meta>,
+}
+
+/// Content types for sampling messages (SEP-1577).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub enum SamplingMessageContent {
+    Text(RawTextContent),
+    Image(RawImageContent),
+    Audio(RawAudioContent),
+    /// Assistant only
+    ToolUse(ToolUseContent),
+    /// User only
+    ToolResult(ToolResultContent),
+}
+
+impl SamplingMessageContent {
+    /// Create a text content
+    pub fn text(text: impl Into<String>) -> Self {
+        Self::Text(RawTextContent {
+            text: text.into(),
+            meta: None,
+        })
+    }
+
+    pub fn tool_use(id: impl Into<String>, name: impl Into<String>, input: JsonObject) -> Self {
+        Self::ToolUse(ToolUseContent::new(id, name, input))
+    }
+
+    pub fn tool_result(tool_use_id: impl Into<String>, content: Vec<Content>) -> Self {
+        Self::ToolResult(ToolResultContent::new(tool_use_id, content))
+    }
+}
+
+impl SamplingMessage {
+    pub fn new(role: Role, content: impl Into<SamplingMessageContent>) -> Self {
+        Self {
+            role,
+            content: SamplingContent::Single(content.into()),
+            meta: None,
+        }
+    }
+
+    pub fn new_multiple(role: Role, contents: Vec<SamplingMessageContent>) -> Self {
+        Self {
+            role,
+            content: SamplingContent::Multiple(contents),
+            meta: None,
+        }
+    }
+
+    pub fn user_text(text: impl Into<String>) -> Self {
+        Self::new(Role::User, SamplingMessageContent::text(text))
+    }
+
+    pub fn assistant_text(text: impl Into<String>) -> Self {
+        Self::new(Role::Assistant, SamplingMessageContent::text(text))
+    }
+
+    pub fn user_tool_result(tool_use_id: impl Into<String>, content: Vec<Content>) -> Self {
+        Self::new(
+            Role::User,
+            SamplingMessageContent::tool_result(tool_use_id, content),
+        )
+    }
+
+    pub fn assistant_tool_use(
+        id: impl Into<String>,
+        name: impl Into<String>,
+        input: JsonObject,
+    ) -> Self {
+        Self::new(
+            Role::Assistant,
+            SamplingMessageContent::tool_use(id, name, input),
+        )
+    }
+}
+
+// Conversion from RawTextContent to SamplingMessageContent
+impl From<RawTextContent> for SamplingMessageContent {
+    fn from(text: RawTextContent) -> Self {
+        SamplingMessageContent::Text(text)
+    }
+}
+
+// Conversion from String to SamplingMessageContent (as text)
+impl From<String> for SamplingMessageContent {
+    fn from(text: String) -> Self {
+        SamplingMessageContent::text(text)
+    }
+}
+
+impl From<&str> for SamplingMessageContent {
+    fn from(text: &str) -> Self {
+        SamplingMessageContent::text(text)
+    }
 }
 
 /// Specifies how much context should be included in sampling requests.
@@ -1267,20 +1506,26 @@ pub struct CreateMessageRequestParams {
     /// System prompt to guide the model's behavior
     #[serde(skip_serializing_if = "Option::is_none")]
     pub system_prompt: Option<String>,
-    /// How much context to include from MCP servers
+    /// Context inclusion (soft-deprecated)
     #[serde(skip_serializing_if = "Option::is_none")]
     pub include_context: Option<ContextInclusion>,
-    /// Temperature for controlling randomness (0.0 to 1.0)
+    /// Temperature (0.0 to 1.0)
     #[serde(skip_serializing_if = "Option::is_none")]
     pub temperature: Option<f32>,
-    /// Maximum number of tokens to generate
+    /// Max tokens to generate
     pub max_tokens: u32,
-    /// Sequences that should stop generation
+    /// Stop sequences
     #[serde(skip_serializing_if = "Option::is_none")]
     pub stop_sequences: Option<Vec<String>>,
-    /// Additional metadata for the request
+    /// Request metadata
     #[serde(skip_serializing_if = "Option::is_none")]
     pub metadata: Option<Value>,
+    /// Tools for the model (SEP-1577)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tools: Option<Vec<Tool>>,
+    /// Tool selection config (SEP-1577)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_choice: Option<ToolChoice>,
 }
 
 impl RequestParamsMeta for CreateMessageRequestParams {
@@ -1926,20 +2171,15 @@ pub type CallToolRequestParam = CallToolRequestParams;
 /// Request to call a specific tool
 pub type CallToolRequest = Request<CallToolRequestMethod, CallToolRequestParams>;
 
-/// The result of a sampling/createMessage request containing the generated response.
-///
-/// This structure contains the generated message along with metadata about
-/// how the generation was performed and why it stopped.
+/// Result of sampling/createMessage (SEP-1577).
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 #[serde(rename_all = "camelCase")]
 #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
 pub struct CreateMessageResult {
-    /// The identifier of the model that generated the response
     pub model: String,
-    /// The reason why generation stopped (e.g., "endTurn", "maxTokens")
+    /// Why generation stopped: "endTurn", "stopSequence", "maxTokens", "toolUse"
     #[serde(skip_serializing_if = "Option::is_none")]
     pub stop_reason: Option<String>,
-    /// The generated message with role and content
     #[serde(flatten)]
     pub message: SamplingMessage,
 }
@@ -1948,6 +2188,7 @@ impl CreateMessageResult {
     pub const STOP_REASON_END_TURN: &str = "endTurn";
     pub const STOP_REASON_END_SEQUENCE: &str = "stopSequence";
     pub const STOP_REASON_END_MAX_TOKEN: &str = "maxTokens";
+    pub const STOP_REASON_TOOL_USE: &str = "toolUse";
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
@@ -2476,7 +2717,10 @@ mod tests {
                 ..
             }) => {
                 assert_eq!(capabilities.roots.unwrap().list_changed, Some(true));
-                assert_eq!(capabilities.sampling.unwrap().len(), 0);
+                // Empty sampling capability (no tools or context sub-capabilities)
+                let sampling = capabilities.sampling.unwrap();
+                assert_eq!(sampling.tools, None);
+                assert_eq!(sampling.context, None);
                 assert_eq!(client_info.name, "ExampleClient");
                 assert_eq!(client_info.version, "1.0.0");
             }