fix(tasks): #626 model task capabilities correctly (#627)

crates/rmcp/src/model.rs

@@ -2089,6 +2089,7 @@ macro_rules! ts_union {
     (@declare_end $U:ident { $($declared:tt)* }) => {
         #[derive(Debug, Serialize, Deserialize, Clone)]
         #[serde(untagged)]
+        #[allow(clippy::large_enum_variant)]
         #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
         pub enum $U {
             $($declared)*

crates/rmcp/src/model/capabilities.rs

@@ -40,24 +40,121 @@ pub struct RootsCapabilities {
     pub list_changed: Option<bool>,
 }
 
-/// Task capability negotiation for SEP-1686.
+/// Task capabilities shared by client and server.
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
 #[serde(rename_all = "camelCase")]
 #[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
 pub struct TasksCapability {
-    /// Map of request category (e.g. "tools.call") to a boolean indicating support.
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub requests: Option<TaskRequestMap>,
-    /// Whether the receiver supports `tasks/list`.
+    pub requests: Option<TaskRequestsCapability>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub list: Option<bool>,
-    /// Whether the receiver supports `tasks/cancel`.
+    pub list: Option<JsonObject>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub cancel: Option<bool>,
+    pub cancel: Option<JsonObject>,
 }
 
-/// A convenience alias for describing per-request task support.
-pub type TaskRequestMap = BTreeMap<String, bool>;
+/// Request types that support task-augmented execution.
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
+#[serde(rename_all = "camelCase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct TaskRequestsCapability {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub sampling: Option<SamplingTaskCapability>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub elicitation: Option<ElicitationTaskCapability>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tools: Option<ToolsTaskCapability>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
+#[serde(rename_all = "camelCase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct SamplingTaskCapability {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub create_message: Option<JsonObject>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
+#[serde(rename_all = "camelCase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct ElicitationTaskCapability {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub create: Option<JsonObject>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
+#[serde(rename_all = "camelCase")]
+#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
+pub struct ToolsTaskCapability {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub call: Option<JsonObject>,
+}
+
+impl TasksCapability {
+    /// Default client tasks capability with sampling and elicitation support.
+    pub fn client_default() -> Self {
+        Self {
+            list: Some(JsonObject::new()),
+            cancel: Some(JsonObject::new()),
+            requests: Some(TaskRequestsCapability {
+                sampling: Some(SamplingTaskCapability {
+                    create_message: Some(JsonObject::new()),
+                }),
+                elicitation: Some(ElicitationTaskCapability {
+                    create: Some(JsonObject::new()),
+                }),
+                tools: None,
+            }),
+        }
+    }
+
+    /// Default server tasks capability with tools/call support.
+    pub fn server_default() -> Self {
+        Self {
+            list: Some(JsonObject::new()),
+            cancel: Some(JsonObject::new()),
+            requests: Some(TaskRequestsCapability {
+                sampling: None,
+                elicitation: None,
+                tools: Some(ToolsTaskCapability {
+                    call: Some(JsonObject::new()),
+                }),
+            }),
+        }
+    }
+
+    pub fn supports_list(&self) -> bool {
+        self.list.is_some()
+    }
+
+    pub fn supports_cancel(&self) -> bool {
+        self.cancel.is_some()
+    }
+
+    pub fn supports_tools_call(&self) -> bool {
+        self.requests
+            .as_ref()
+            .and_then(|r| r.tools.as_ref())
+            .and_then(|t| t.call.as_ref())
+            .is_some()
+    }
+
+    pub fn supports_sampling_create_message(&self) -> bool {
+        self.requests
+            .as_ref()
+            .and_then(|r| r.sampling.as_ref())
+            .and_then(|s| s.create_message.as_ref())
+            .is_some()
+    }
+
+    pub fn supports_elicitation_create(&self) -> bool {
+        self.requests
+            .as_ref()
+            .and_then(|r| r.elicitation.as_ref())
+            .and_then(|e| e.create.as_ref())
+            .is_some()
+    }
+}
 
 /// Capability for handling elicitation requests from servers.
 ///
@@ -368,4 +465,67 @@ mod test {
             })
         );
     }
+
+    #[test]
+    fn test_task_capabilities_deserialization() {
+        // Test deserializing from the MCP spec format
+        let json = serde_json::json!({
+            "list": {},
+            "cancel": {},
+            "requests": {
+                "tools": { "call": {} }
+            }
+        });
+
+        let tasks: TasksCapability = serde_json::from_value(json).unwrap();
+        assert!(tasks.list.is_some());
+        assert!(tasks.cancel.is_some());
+        assert!(tasks.requests.is_some());
+        let requests = tasks.requests.unwrap();
+        assert!(requests.tools.is_some());
+        assert!(requests.tools.unwrap().call.is_some());
+    }
+
+    #[test]
+    fn test_tasks_capability_client_default() {
+        let tasks = TasksCapability::client_default();
+
+        // Verify structure
+        assert!(tasks.supports_list());
+        assert!(tasks.supports_cancel());
+        assert!(tasks.supports_sampling_create_message());
+        assert!(tasks.supports_elicitation_create());
+        assert!(!tasks.supports_tools_call());
+
+        // Verify serialization matches expected format
+        let json = serde_json::to_value(&tasks).unwrap();
+        assert_eq!(json["list"], serde_json::json!({}));
+        assert_eq!(json["cancel"], serde_json::json!({}));
+        assert_eq!(
+            json["requests"]["sampling"]["createMessage"],
+            serde_json::json!({})
+        );
+        assert_eq!(
+            json["requests"]["elicitation"]["create"],
+            serde_json::json!({})
+        );
+    }
+
+    #[test]
+    fn test_tasks_capability_server_default() {
+        let tasks = TasksCapability::server_default();
+
+        // Verify structure
+        assert!(tasks.supports_list());
+        assert!(tasks.supports_cancel());
+        assert!(tasks.supports_tools_call());
+        assert!(!tasks.supports_sampling_create_message());
+        assert!(!tasks.supports_elicitation_create());
+
+        // Verify serialization matches expected format
+        let json = serde_json::to_value(&tasks).unwrap();
+        assert_eq!(json["list"], serde_json::json!({}));
+        assert_eq!(json["cancel"], serde_json::json!({}));
+        assert_eq!(json["requests"]["tools"]["call"], serde_json::json!({}));
+    }
 }