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

2026-01-23 13:03:31 -05:00 · 2026-01-23 13:03:31 -05:00 · 613eafbda8
commit 613eafbda8
parent e623f2acab
6 changed files with 510 additions and 61 deletions
--- a/crates/rmcp/src/model.rs
+++ b/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)*
--- a/crates/rmcp/src/model/capabilities.rs
+++ b/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!({}));
+    }
 }
--- a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json
+++ b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema.json
@ -519,6 +519,18 @@
        }
      }
    },
+    "ElicitationTaskCapability": {
+      "type": "object",
+      "properties": {
+        "create": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "EmptyObject": {
      "description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
      "type": "object",
@ -1718,6 +1730,18 @@
      "format": "const",
      "const": "notifications/roots/list_changed"
    },
+    "SamplingTaskCapability": {
+      "type": "object",
+      "properties": {
+        "createMessage": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "SetLevelRequestMethod": {
      "type": "string",
      "format": "const",
@ -1774,33 +1798,81 @@
        "uri"
      ]
    },
+    "TaskRequestsCapability": {
+      "description": "Request types that support task-augmented execution.",
+      "type": "object",
+      "properties": {
+        "elicitation": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ElicitationTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "sampling": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/SamplingTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "tools": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ToolsTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        }
+      }
+    },
    "TasksCapability": {
-      "description": "Task capability negotiation for SEP-1686.",
+      "description": "Task capabilities shared by client and server.",
      "type": "object",
      "properties": {
        "cancel": {
-          "description": "Whether the receiver supports `tasks/cancel`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "list": {
-          "description": "Whether the receiver supports `tasks/list`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "requests": {
-          "description": "Map of request category (e.g. \"tools.call\") to a boolean indicating support.",
          "type": [
            "object",
            "null"
          ],
-          "additionalProperties": {
-            "type": "boolean"
-          }
+          "additionalProperties": true
+        },
+        "list": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        },
+        "requests": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/TaskRequestsCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        }
+      }
+    },
+    "ToolsTaskCapability": {
+      "type": "object",
+      "properties": {
+        "call": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
        }
      }
    },
--- a/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json
+++ b/crates/rmcp/tests/test_message_schema/client_json_rpc_message_schema_current.json
@ -519,6 +519,18 @@
        }
      }
    },
+    "ElicitationTaskCapability": {
+      "type": "object",
+      "properties": {
+        "create": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "EmptyObject": {
      "description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
      "type": "object",
@ -1718,6 +1730,18 @@
      "format": "const",
      "const": "notifications/roots/list_changed"
    },
+    "SamplingTaskCapability": {
+      "type": "object",
+      "properties": {
+        "createMessage": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "SetLevelRequestMethod": {
      "type": "string",
      "format": "const",
@ -1774,33 +1798,81 @@
        "uri"
      ]
    },
+    "TaskRequestsCapability": {
+      "description": "Request types that support task-augmented execution.",
+      "type": "object",
+      "properties": {
+        "elicitation": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ElicitationTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "sampling": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/SamplingTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "tools": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ToolsTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        }
+      }
+    },
    "TasksCapability": {
-      "description": "Task capability negotiation for SEP-1686.",
+      "description": "Task capabilities shared by client and server.",
      "type": "object",
      "properties": {
        "cancel": {
-          "description": "Whether the receiver supports `tasks/cancel`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "list": {
-          "description": "Whether the receiver supports `tasks/list`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "requests": {
-          "description": "Map of request category (e.g. \"tools.call\") to a boolean indicating support.",
          "type": [
            "object",
            "null"
          ],
-          "additionalProperties": {
-            "type": "boolean"
-          }
+          "additionalProperties": true
+        },
+        "list": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        },
+        "requests": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/TaskRequestsCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        }
+      }
+    },
+    "ToolsTaskCapability": {
+      "type": "object",
+      "properties": {
+        "call": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
        }
      }
    },
--- a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json
+++ b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema.json
@ -763,6 +763,18 @@
        "properties"
      ]
    },
+    "ElicitationTaskCapability": {
+      "type": "object",
+      "properties": {
+        "create": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "EmptyObject": {
      "description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
      "type": "object",
@ -2322,6 +2334,18 @@
        "content"
      ]
    },
+    "SamplingTaskCapability": {
+      "type": "object",
+      "properties": {
+        "createMessage": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "ServerCapabilities": {
      "title": "Builder",
      "description": "```rust\n# use rmcp::model::ServerCapabilities;\nlet cap = ServerCapabilities::builder()\n    .enable_logging()\n    .enable_experimental()\n    .enable_prompts()\n    .enable_resources()\n    .enable_tools()\n    .enable_tool_list_changed()\n    .build();\n```",
@ -2605,6 +2629,42 @@
        "createdAt"
      ]
    },
+    "TaskRequestsCapability": {
+      "description": "Request types that support task-augmented execution.",
+      "type": "object",
+      "properties": {
+        "elicitation": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ElicitationTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "sampling": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/SamplingTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "tools": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ToolsTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        }
+      }
+    },
    "TaskResult": {
      "description": "Final result for a succeeded task (returned from `tasks/result`).",
      "type": "object",
@ -2660,32 +2720,32 @@
      ]
    },
    "TasksCapability": {
-      "description": "Task capability negotiation for SEP-1686.",
+      "description": "Task capabilities shared by client and server.",
      "type": "object",
      "properties": {
        "cancel": {
-          "description": "Whether the receiver supports `tasks/cancel`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "list": {
-          "description": "Whether the receiver supports `tasks/list`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "requests": {
-          "description": "Map of request category (e.g. \"tools.call\") to a boolean indicating support.",
          "type": [
            "object",
            "null"
          ],
-          "additionalProperties": {
-            "type": "boolean"
-          }
+          "additionalProperties": true
+        },
+        "list": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        },
+        "requests": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/TaskRequestsCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
        }
      }
    },
@ -2921,6 +2981,18 @@
        }
      }
    },
+    "ToolsTaskCapability": {
+      "type": "object",
+      "properties": {
+        "call": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "UntitledItems": {
      "description": "Items for untitled multi-select options",
      "type": "object",
--- a/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json
+++ b/crates/rmcp/tests/test_message_schema/server_json_rpc_message_schema_current.json
@ -763,6 +763,18 @@
        "properties"
      ]
    },
+    "ElicitationTaskCapability": {
+      "type": "object",
+      "properties": {
+        "create": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "EmptyObject": {
      "description": "This is commonly used for representing empty objects in MCP messages.\n\nwithout returning any specific data.",
      "type": "object",
@ -2322,6 +2334,18 @@
        "content"
      ]
    },
+    "SamplingTaskCapability": {
+      "type": "object",
+      "properties": {
+        "createMessage": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "ServerCapabilities": {
      "title": "Builder",
      "description": "```rust\n# use rmcp::model::ServerCapabilities;\nlet cap = ServerCapabilities::builder()\n    .enable_logging()\n    .enable_experimental()\n    .enable_prompts()\n    .enable_resources()\n    .enable_tools()\n    .enable_tool_list_changed()\n    .build();\n```",
@ -2605,6 +2629,42 @@
        "createdAt"
      ]
    },
+    "TaskRequestsCapability": {
+      "description": "Request types that support task-augmented execution.",
+      "type": "object",
+      "properties": {
+        "elicitation": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ElicitationTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "sampling": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/SamplingTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        },
+        "tools": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/ToolsTaskCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
+        }
+      }
+    },
    "TaskResult": {
      "description": "Final result for a succeeded task (returned from `tasks/result`).",
      "type": "object",
@ -2660,32 +2720,32 @@
      ]
    },
    "TasksCapability": {
-      "description": "Task capability negotiation for SEP-1686.",
+      "description": "Task capabilities shared by client and server.",
      "type": "object",
      "properties": {
        "cancel": {
-          "description": "Whether the receiver supports `tasks/cancel`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "list": {
-          "description": "Whether the receiver supports `tasks/list`.",
-          "type": [
-            "boolean",
-            "null"
-          ]
-        },
-        "requests": {
-          "description": "Map of request category (e.g. \"tools.call\") to a boolean indicating support.",
          "type": [
            "object",
            "null"
          ],
-          "additionalProperties": {
-            "type": "boolean"
-          }
+          "additionalProperties": true
+        },
+        "list": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        },
+        "requests": {
+          "anyOf": [
+            {
+              "$ref": "#/definitions/TaskRequestsCapability"
+            },
+            {
+              "type": "null"
+            }
+          ]
        }
      }
    },
@ -2921,6 +2981,18 @@
        }
      }
    },
+    "ToolsTaskCapability": {
+      "type": "object",
+      "properties": {
+        "call": {
+          "type": [
+            "object",
+            "null"
+          ],
+          "additionalProperties": true
+        }
+      }
+    },
    "UntitledItems": {
      "description": "Items for untitled multi-select options",
      "type": "object",