feat: implement SEP-1577 sampling with tools support

Author: DaleSeo

crates/rmcp/src/model.rs

@@ -1209,18 +1209,292 @@ pub enum Role {
     Assistant,
 }
 
+/// Tool selection mode for sampling requests (SEP-1577).
+///
+/// Controls how the model should handle tool calling during message generation.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub enum ToolChoiceMode {
+    /// Let the model decide whether to call tools
+    Auto,
+    /// The model must call at least one tool
+    Required,
+    /// The model must not call any tools
+    None,
+}
+
+impl Default for ToolChoiceMode {
+    fn default() -> Self {
+        Self::Auto
+    }
+}
+
+/// Tool choice configuration for sampling requests (SEP-1577).
+///
+/// Controls how the model should select and use tools when generating responses.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
+#[serde(rename_all = "camelCase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct ToolChoice {
+    /// The mode for tool selection
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub mode: Option<ToolChoiceMode>,
+}
+
+impl ToolChoice {
+    /// Create a new ToolChoice with auto mode
+    pub fn auto() -> Self {
+        Self {
+            mode: Some(ToolChoiceMode::Auto),
+        }
+    }
+
+    /// Create a new ToolChoice with required mode
+    pub fn required() -> Self {
+        Self {
+            mode: Some(ToolChoiceMode::Required),
+        }
+    }
+
+    /// Create a new ToolChoice with none mode
+    pub fn none() -> Self {
+        Self {
+            mode: Some(ToolChoiceMode::None),
+        }
+    }
+}
+
+/// Content for a sampling message that can be a single item or an array (SEP-1577).
+///
+/// This wrapper type handles the fact that content can be either a single
+/// content block or an array of content blocks, depending on the context.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub enum SamplingContent<T> {
+    /// A single content block
+    Single(T),
+    /// Multiple content blocks
+    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, 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.
+///
+/// Per SEP-1577, the content can be a single content block or an array of
+/// content blocks. For user messages, content can include tool results.
+/// For assistant messages, content can include tool use requests.
 #[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,
+    /// The actual content of the message (text, image, tool use/result, etc.)
+    /// Can be a single content or an array of contents (SEP-1577)
+    pub content: SamplingContent<SamplingMessageContent>,
+}
+
+/// Content types for sampling messages (SEP-1577).
+///
+/// This enum represents all possible content types that can appear in
+/// sampling messages. The appropriate content types depend on the role:
+/// - User messages: text, image, audio, tool_result
+/// - Assistant messages: text, image, audio, tool_use
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub enum SamplingMessageContent {
+    /// Plain text content
+    Text(RawTextContent),
+    /// Image content with base64-encoded data
+    Image(RawImageContent),
+    /// Audio content with base64-encoded data
+    Audio(RawAudioContent),
+    /// A request to use a tool (assistant messages only, SEP-1577)
+    ToolUse(ToolUseContent),
+    /// The result of a tool call (user messages only, SEP-1577)
+    ToolResult(ToolResultContent),
+}
+
+impl SamplingMessageContent {
+    /// Create a text content
+    pub fn text(text: impl Into<String>) -> Self {
+        Self::Text(RawTextContent {
+            text: text.into(),
+            meta: None,
+        })
+    }
+
+    /// Create a tool use content
+    pub fn tool_use(id: impl Into<String>, name: impl Into<String>, input: JsonObject) -> Self {
+        Self::ToolUse(ToolUseContent::new(id, name, input))
+    }
+
+    /// Create a tool result content
+    pub fn tool_result(tool_use_id: impl Into<String>, content: Vec<Content>) -> Self {
+        Self::ToolResult(ToolResultContent::new(tool_use_id, content))
+    }
+}
+
+impl SamplingMessage {
+    /// Create a new sampling message with a single content
+    pub fn new(role: Role, content: impl Into<SamplingMessageContent>) -> Self {
+        Self {
+            role,
+            content: SamplingContent::Single(content.into()),
+        }
+    }
+
+    /// Create a new sampling message with multiple contents
+    pub fn new_multiple(role: Role, contents: Vec<SamplingMessageContent>) -> Self {
+        Self {
+            role,
+            content: SamplingContent::Multiple(contents),
+        }
+    }
+
+    /// Create a user message with text content
+    pub fn user_text(text: impl Into<String>) -> Self {
+        Self::new(Role::User, SamplingMessageContent::text(text))
+    }
+
+    /// Create an assistant message with text content
+    pub fn assistant_text(text: impl Into<String>) -> Self {
+        Self::new(Role::Assistant, SamplingMessageContent::text(text))
+    }
+
+    /// Create a user message with tool result content
+    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),
+        )
+    }
+
+    /// Create an assistant message with tool use 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,7 +1541,9 @@ 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
+    /// How much context to include from MCP servers.
+    /// Note: Values other than "none" are soft-deprecated per SEP-1577
+    /// and require `clientCapabilities.sampling.context` support.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub include_context: Option<ContextInclusion>,
     /// Temperature for controlling randomness (0.0 to 1.0)
@@ -1281,6 +1557,14 @@ pub struct CreateMessageRequestParams {
     /// Additional metadata for the request
     #[serde(skip_serializing_if = "Option::is_none")]
     pub metadata: Option<Value>,
+    /// Tools available for the model to call (SEP-1577).
+    /// Requires `clientCapabilities.sampling.tools` support.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tools: Option<Vec<Tool>>,
+    /// Configuration for tool selection behavior (SEP-1577).
+    /// Requires `clientCapabilities.sampling.tools` support.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_choice: Option<ToolChoice>,
 }
 
 impl RequestParamsMeta for CreateMessageRequestParams {
@@ -1930,13 +2214,17 @@ pub type CallToolRequest = Request<CallToolRequestMethod, CallToolRequestParams>
 ///
 /// This structure contains the generated message along with metadata about
 /// how the generation was performed and why it stopped.
+///
+/// Per SEP-1577, the content can be a single content block or an array of
+/// content blocks when the model uses tools.
 #[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")
+    /// The reason why generation stopped.
+    /// Common values: "endTurn", "stopSequence", "maxTokens", "toolUse"
     #[serde(skip_serializing_if = "Option::is_none")]
     pub stop_reason: Option<String>,
     /// The generated message with role and content
@@ -1945,9 +2233,14 @@ pub struct CreateMessageResult {
 }
 
 impl CreateMessageResult {
+    /// Stop reason: The model naturally ended its turn
     pub const STOP_REASON_END_TURN: &str = "endTurn";
+    /// Stop reason: A stop sequence was encountered
     pub const STOP_REASON_END_SEQUENCE: &str = "stopSequence";
+    /// Stop reason: Maximum token limit reached
     pub const STOP_REASON_END_MAX_TOKEN: &str = "maxTokens";
+    /// Stop reason: The model wants to use a tool (SEP-1577)
+    pub const STOP_REASON_TOOL_USE: &str = "toolUse";
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
@@ -2476,7 +2769,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");
             }