fix(tasks): expose execution.taskSupport on tools (#635)
* fix(tasks): expose execution.taskSupport on tools * feat: implement taskSupport validation on server
This commit is contained in:
Luca Chang
GitHub
parent
1794fe1548
commit
df6c3f0665
10 changed files with 539 additions and 2 deletions
|
|
@ -89,12 +89,21 @@ pub struct ToolAttribute {
|
|||
pub output_schema: Option<Expr>,
|
||||
/// Optional additional tool information.
|
||||
pub annotations: Option<ToolAnnotationsAttribute>,
|
||||
/// Execution-related configuration including task support.
|
||||
pub execution: Option<ToolExecutionAttribute>,
|
||||
/// Optional icons for the tool
|
||||
pub icons: Option<Expr>,
|
||||
/// Optional metadata for the tool
|
||||
pub meta: Option<Expr>,
|
||||
}
|
||||
|
||||
#[derive(FromMeta, Debug, Default)]
|
||||
#[darling(default)]
|
||||
pub struct ToolExecutionAttribute {
|
||||
/// Task support mode: "forbidden", "optional", or "required"
|
||||
pub task_support: Option<String>,
|
||||
}
|
||||
|
||||
pub struct ResolvedToolAttribute {
|
||||
pub name: String,
|
||||
pub title: Option<String>,
|
||||
|
|
@ -102,6 +111,7 @@ pub struct ResolvedToolAttribute {
|
|||
pub input_schema: Expr,
|
||||
pub output_schema: Option<Expr>,
|
||||
pub annotations: Expr,
|
||||
pub execution: Expr,
|
||||
pub icons: Option<Expr>,
|
||||
pub meta: Option<Expr>,
|
||||
}
|
||||
|
|
@ -115,6 +125,7 @@ impl ResolvedToolAttribute {
|
|||
input_schema,
|
||||
output_schema,
|
||||
annotations,
|
||||
execution,
|
||||
icons,
|
||||
meta,
|
||||
} = self;
|
||||
|
|
@ -155,6 +166,7 @@ impl ResolvedToolAttribute {
|
|||
input_schema: #input_schema,
|
||||
output_schema: #output_schema,
|
||||
annotations: #annotations,
|
||||
execution: #execution,
|
||||
icons: #icons,
|
||||
meta: #meta,
|
||||
}
|
||||
|
|
@ -263,6 +275,38 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
|
|||
} else {
|
||||
none_expr()?
|
||||
};
|
||||
let execution_expr = if let Some(execution) = attribute.execution {
|
||||
let ToolExecutionAttribute { task_support } = execution;
|
||||
|
||||
let task_support_expr = if let Some(ts) = task_support {
|
||||
let ts_ident = match ts.as_str() {
|
||||
"forbidden" => quote! { rmcp::model::TaskSupport::Forbidden },
|
||||
"optional" => quote! { rmcp::model::TaskSupport::Optional },
|
||||
"required" => quote! { rmcp::model::TaskSupport::Required },
|
||||
_ => {
|
||||
return Err(syn::Error::new(
|
||||
Span::call_site(),
|
||||
format!(
|
||||
"Invalid task_support value '{}'. Expected 'forbidden', 'optional', or 'required'",
|
||||
ts
|
||||
),
|
||||
));
|
||||
}
|
||||
};
|
||||
quote! { Some(#ts_ident) }
|
||||
} else {
|
||||
quote! { None }
|
||||
};
|
||||
|
||||
let token_stream = quote! {
|
||||
Some(rmcp::model::ToolExecution {
|
||||
task_support: #task_support_expr,
|
||||
})
|
||||
};
|
||||
syn::parse2::<Expr>(token_stream)?
|
||||
} else {
|
||||
none_expr()?
|
||||
};
|
||||
// Handle output_schema - either explicit or generated from return type
|
||||
let output_schema_expr = attribute.output_schema.or_else(|| {
|
||||
// Try to generate schema from return type
|
||||
|
|
@ -286,6 +330,7 @@ pub fn tool(attr: TokenStream, input: TokenStream) -> syn::Result<TokenStream> {
|
|||
input_schema: input_schema_expr,
|
||||
output_schema: output_schema_expr,
|
||||
annotations: annotations_expr,
|
||||
execution: execution_expr,
|
||||
title: attribute.title,
|
||||
icons: attribute.icons,
|
||||
meta: attribute.meta,
|
||||
|
|
|
|||
|
|
@ -56,9 +56,18 @@ pub fn tool_handler(attr: TokenStream, input: TokenStream) -> syn::Result<TokenS
|
|||
})
|
||||
}
|
||||
};
|
||||
|
||||
let get_tool_fn = quote! {
|
||||
fn get_tool(&self, name: &str) -> Option<rmcp::model::Tool> {
|
||||
#router.get(name).cloned()
|
||||
}
|
||||
};
|
||||
|
||||
let tool_call_fn = syn::parse2::<ImplItem>(tool_call_fn)?;
|
||||
let tool_list_fn = syn::parse2::<ImplItem>(tool_list_fn)?;
|
||||
let get_tool_fn = syn::parse2::<ImplItem>(get_tool_fn)?;
|
||||
item_impl.items.push(tool_call_fn);
|
||||
item_impl.items.push(tool_list_fn);
|
||||
item_impl.items.push(get_tool_fn);
|
||||
Ok(item_impl.into_token_stream())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use std::sync::Arc;
|
|||
|
||||
use crate::{
|
||||
error::ErrorData as McpError,
|
||||
model::*,
|
||||
model::{TaskSupport, *},
|
||||
service::{NotificationContext, RequestContext, RoleServer, Service, ServiceRole},
|
||||
};
|
||||
|
||||
|
|
@ -65,7 +65,32 @@ impl<H: ServerHandler> Service<RoleServer> for H {
|
|||
.await
|
||||
.map(ServerResult::empty),
|
||||
ClientRequest::CallToolRequest(request) => {
|
||||
if request.params.task.is_some() {
|
||||
let is_task = request.params.task.is_some();
|
||||
|
||||
// Validate task support mode per MCP specification
|
||||
if let Some(tool) = self.get_tool(&request.params.name) {
|
||||
match (tool.task_support(), is_task) {
|
||||
// If taskSupport is "required", clients MUST invoke the tool as a task.
|
||||
// Servers MUST return a -32601 (Method not found) error if they don't.
|
||||
(TaskSupport::Required, false) => {
|
||||
return Err(McpError::new(
|
||||
ErrorCode::METHOD_NOT_FOUND,
|
||||
"Tool requires task-based invocation",
|
||||
None,
|
||||
));
|
||||
}
|
||||
// If taskSupport is "forbidden" (default), clients MUST NOT invoke as a task.
|
||||
(TaskSupport::Forbidden, true) => {
|
||||
return Err(McpError::invalid_params(
|
||||
"Tool does not support task-based invocation",
|
||||
None,
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if is_task {
|
||||
tracing::info!("Enqueueing task for tool call: {}", request.params.name);
|
||||
self.enqueue_task(request.params, context.clone())
|
||||
.await
|
||||
|
|
@ -241,6 +266,13 @@ pub trait ServerHandler: Sized + Send + Sync + 'static {
|
|||
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(ListToolsResult::default()))
|
||||
}
|
||||
/// Get a tool definition by name.
|
||||
///
|
||||
/// The default implementation returns `None`, which bypasses validation.
|
||||
/// When using `#[tool_handler]`, this method is automatically implemented.
|
||||
fn get_tool(&self, _name: &str) -> Option<Tool> {
|
||||
None
|
||||
}
|
||||
fn on_custom_request(
|
||||
&self,
|
||||
request: CustomRequest,
|
||||
|
|
@ -445,6 +477,10 @@ macro_rules! impl_server_handler_for_wrapper {
|
|||
(**self).list_tools(request, context)
|
||||
}
|
||||
|
||||
fn get_tool(&self, name: &str) -> Option<Tool> {
|
||||
(**self).get_tool(name)
|
||||
}
|
||||
|
||||
fn on_custom_request(
|
||||
&self,
|
||||
request: CustomRequest,
|
||||
|
|
|
|||
|
|
@ -254,6 +254,13 @@ where
|
|||
pub fn list_all(&self) -> Vec<crate::model::Tool> {
|
||||
self.map.values().map(|item| item.attr.clone()).collect()
|
||||
}
|
||||
|
||||
/// Get a tool definition by name.
|
||||
///
|
||||
/// Returns the tool if found, or `None` if no tool with the given name exists.
|
||||
pub fn get(&self, name: &str) -> Option<&crate::model::Tool> {
|
||||
self.map.get(name).map(|r| &r.attr)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S> std::ops::Add<ToolRouter<S>> for ToolRouter<S>
|
||||
|
|
|
|||
|
|
@ -29,6 +29,9 @@ pub struct Tool {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
/// Optional additional tool information.
|
||||
pub annotations: Option<ToolAnnotations>,
|
||||
/// Execution-related configuration including task support mode.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub execution: Option<ToolExecution>,
|
||||
/// Optional list of icons for the tool
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub icons: Option<Vec<Icon>>,
|
||||
|
|
@ -37,6 +40,55 @@ pub struct Tool {
|
|||
pub meta: Option<Meta>,
|
||||
}
|
||||
|
||||
/// Per-tool task support mode as defined in the MCP specification.
|
||||
///
|
||||
/// This enum indicates whether a tool supports task-based invocation,
|
||||
/// allowing clients to know how to properly call the tool.
|
||||
///
|
||||
/// See [Tool-Level Negotiation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks#tool-level-negotiation).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
|
||||
pub enum TaskSupport {
|
||||
/// Clients MUST NOT invoke this tool as a task (default behavior).
|
||||
#[default]
|
||||
Forbidden,
|
||||
/// Clients MAY invoke this tool as either a task or a normal call.
|
||||
Optional,
|
||||
/// Clients MUST invoke this tool as a task.
|
||||
Required,
|
||||
}
|
||||
|
||||
/// Execution-related configuration for a tool.
|
||||
///
|
||||
/// This struct contains settings that control how a tool should be executed,
|
||||
/// including task support configuration.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
|
||||
pub struct ToolExecution {
|
||||
/// Indicates whether this tool supports task-based invocation.
|
||||
///
|
||||
/// When not present or set to `Forbidden`, clients MUST NOT invoke this tool as a task.
|
||||
/// When set to `Optional`, clients MAY invoke this tool as a task or normal call.
|
||||
/// When set to `Required`, clients MUST invoke this tool as a task.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub task_support: Option<TaskSupport>,
|
||||
}
|
||||
|
||||
impl ToolExecution {
|
||||
/// Create a new empty ToolExecution configuration.
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Set the task support mode.
|
||||
pub fn with_task_support(mut self, task_support: TaskSupport) -> Self {
|
||||
self.task_support = Some(task_support);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Additional properties describing a Tool to clients.
|
||||
///
|
||||
/// NOTE: all properties in ToolAnnotations are **hints**.
|
||||
|
|
@ -152,6 +204,7 @@ impl Tool {
|
|||
input_schema: input_schema.into(),
|
||||
output_schema: None,
|
||||
annotations: None,
|
||||
execution: None,
|
||||
icons: None,
|
||||
meta: None,
|
||||
}
|
||||
|
|
@ -164,6 +217,24 @@ impl Tool {
|
|||
}
|
||||
}
|
||||
|
||||
/// Set the execution configuration for this tool.
|
||||
pub fn with_execution(self, execution: ToolExecution) -> Self {
|
||||
Tool {
|
||||
execution: Some(execution),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the task support mode for this tool.
|
||||
///
|
||||
/// Returns `TaskSupport::Forbidden` if not explicitly set.
|
||||
pub fn task_support(&self) -> TaskSupport {
|
||||
self.execution
|
||||
.as_ref()
|
||||
.and_then(|e| e.task_support)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Set the output schema using a type that implements JsonSchema
|
||||
///
|
||||
/// # Panics
|
||||
|
|
|
|||
|
|
@ -19,6 +19,9 @@
|
|||
],
|
||||
"additionalProperties": false,
|
||||
"$schema": "http://json-schema.org/draft-07/schema#"
|
||||
},
|
||||
"execution": {
|
||||
"taskSupport": "optional"
|
||||
}
|
||||
}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -2731,6 +2731,26 @@
|
|||
}
|
||||
]
|
||||
},
|
||||
"TaskSupport": {
|
||||
"description": "Per-tool task support mode as defined in the MCP specification.\n\nThis enum indicates whether a tool supports task-based invocation,\nallowing clients to know how to properly call the tool.\n\nSee [Tool-Level Negotiation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks#tool-level-negotiation).",
|
||||
"oneOf": [
|
||||
{
|
||||
"description": "Clients MUST NOT invoke this tool as a task (default behavior).",
|
||||
"type": "string",
|
||||
"const": "forbidden"
|
||||
},
|
||||
{
|
||||
"description": "Clients MAY invoke this tool as either a task or a normal call.",
|
||||
"type": "string",
|
||||
"const": "optional"
|
||||
},
|
||||
{
|
||||
"description": "Clients MUST invoke this tool as a task.",
|
||||
"type": "string",
|
||||
"const": "required"
|
||||
}
|
||||
]
|
||||
},
|
||||
"TasksCapability": {
|
||||
"description": "Task capabilities shared by client and server.",
|
||||
"type": "object",
|
||||
|
|
@ -2896,6 +2916,17 @@
|
|||
"null"
|
||||
]
|
||||
},
|
||||
"execution": {
|
||||
"description": "Execution-related configuration including task support mode.",
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/definitions/ToolExecution"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"icons": {
|
||||
"description": "Optional list of icons for the tool",
|
||||
"type": [
|
||||
|
|
@ -2977,6 +3008,23 @@
|
|||
}
|
||||
}
|
||||
},
|
||||
"ToolExecution": {
|
||||
"description": "Execution-related configuration for a tool.\n\nThis struct contains settings that control how a tool should be executed,\nincluding task support configuration.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"taskSupport": {
|
||||
"description": "Indicates whether this tool supports task-based invocation.\n\nWhen not present or set to `Forbidden`, clients MUST NOT invoke this tool as a task.\nWhen set to `Optional`, clients MAY invoke this tool as a task or normal call.\nWhen set to `Required`, clients MUST invoke this tool as a task.",
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/definitions/TaskSupport"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
"ToolListChangedNotificationMethod": {
|
||||
"type": "string",
|
||||
"format": "const",
|
||||
|
|
|
|||
|
|
@ -2731,6 +2731,26 @@
|
|||
}
|
||||
]
|
||||
},
|
||||
"TaskSupport": {
|
||||
"description": "Per-tool task support mode as defined in the MCP specification.\n\nThis enum indicates whether a tool supports task-based invocation,\nallowing clients to know how to properly call the tool.\n\nSee [Tool-Level Negotiation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks#tool-level-negotiation).",
|
||||
"oneOf": [
|
||||
{
|
||||
"description": "Clients MUST NOT invoke this tool as a task (default behavior).",
|
||||
"type": "string",
|
||||
"const": "forbidden"
|
||||
},
|
||||
{
|
||||
"description": "Clients MAY invoke this tool as either a task or a normal call.",
|
||||
"type": "string",
|
||||
"const": "optional"
|
||||
},
|
||||
{
|
||||
"description": "Clients MUST invoke this tool as a task.",
|
||||
"type": "string",
|
||||
"const": "required"
|
||||
}
|
||||
]
|
||||
},
|
||||
"TasksCapability": {
|
||||
"description": "Task capabilities shared by client and server.",
|
||||
"type": "object",
|
||||
|
|
@ -2896,6 +2916,17 @@
|
|||
"null"
|
||||
]
|
||||
},
|
||||
"execution": {
|
||||
"description": "Execution-related configuration including task support mode.",
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/definitions/ToolExecution"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"icons": {
|
||||
"description": "Optional list of icons for the tool",
|
||||
"type": [
|
||||
|
|
@ -2977,6 +3008,23 @@
|
|||
}
|
||||
}
|
||||
},
|
||||
"ToolExecution": {
|
||||
"description": "Execution-related configuration for a tool.\n\nThis struct contains settings that control how a tool should be executed,\nincluding task support configuration.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"taskSupport": {
|
||||
"description": "Indicates whether this tool supports task-based invocation.\n\nWhen not present or set to `Forbidden`, clients MUST NOT invoke this tool as a task.\nWhen set to `Optional`, clients MAY invoke this tool as a task or normal call.\nWhen set to `Required`, clients MUST invoke this tool as a task.",
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/definitions/TaskSupport"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
"ToolListChangedNotificationMethod": {
|
||||
"type": "string",
|
||||
"format": "const",
|
||||
|
|
|
|||
269
crates/rmcp/tests/test_task_support_validation.rs
Normal file
269
crates/rmcp/tests/test_task_support_validation.rs
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
//! Tests for task support validation in tool calls.
|
||||
//!
|
||||
//! Verifies that the server correctly validates `execution.taskSupport` settings
|
||||
//! per the MCP specification:
|
||||
//! - `Required`: MUST be invoked as a task, returns -32601 otherwise
|
||||
//! - `Forbidden`: MUST NOT be invoked as a task, returns error otherwise
|
||||
//! - `Optional`: MAY be invoked either way
|
||||
|
||||
use rmcp::{
|
||||
ClientHandler, ServerHandler, ServiceError, ServiceExt,
|
||||
handler::server::router::tool::ToolRouter,
|
||||
model::{CallToolRequestParams, ClientInfo, ErrorCode, JsonObject},
|
||||
tool, tool_handler, tool_router,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
/// Server with tools having different task support modes.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TaskSupportTestServer {
|
||||
tool_router: ToolRouter<Self>,
|
||||
}
|
||||
|
||||
impl Default for TaskSupportTestServer {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl TaskSupportTestServer {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tool_router: Self::tool_router(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tool_router]
|
||||
impl TaskSupportTestServer {
|
||||
#[tool(
|
||||
description = "Tool that requires task-based invocation",
|
||||
execution(task_support = "required")
|
||||
)]
|
||||
async fn required_task_tool(&self) -> String {
|
||||
"required task executed".to_string()
|
||||
}
|
||||
|
||||
#[tool(
|
||||
description = "Tool that forbids task-based invocation",
|
||||
execution(task_support = "forbidden")
|
||||
)]
|
||||
async fn forbidden_task_tool(&self) -> String {
|
||||
"forbidden task executed".to_string()
|
||||
}
|
||||
|
||||
#[tool(
|
||||
description = "Tool that optionally supports task-based invocation",
|
||||
execution(task_support = "optional")
|
||||
)]
|
||||
async fn optional_task_tool(&self) -> String {
|
||||
"optional task executed".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[tool_handler]
|
||||
impl ServerHandler for TaskSupportTestServer {}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct DummyClientHandler {}
|
||||
|
||||
impl ClientHandler for DummyClientHandler {
|
||||
fn get_info(&self) -> ClientInfo {
|
||||
ClientInfo::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper to create a task object for tool calls
|
||||
fn make_task() -> Option<JsonObject> {
|
||||
Some(json!({}).as_object().unwrap().clone())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_required_task_tool_without_task_returns_method_not_found() -> anyhow::Result<()> {
|
||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||
|
||||
let server = TaskSupportTestServer::new();
|
||||
let server_handle = tokio::spawn(async move {
|
||||
server.serve(server_transport).await?.waiting().await?;
|
||||
anyhow::Ok(())
|
||||
});
|
||||
|
||||
let client_handler = DummyClientHandler::default();
|
||||
let client = client_handler.serve(client_transport).await?;
|
||||
|
||||
// Call the task-required tool without a task - should fail with -32601
|
||||
let result = client
|
||||
.call_tool(CallToolRequestParams {
|
||||
meta: None,
|
||||
name: "required_task_tool".into(),
|
||||
arguments: None,
|
||||
task: None, // No task provided!
|
||||
})
|
||||
.await;
|
||||
|
||||
// Should be an error with code -32601 (METHOD_NOT_FOUND)
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Expected error for required task tool without task"
|
||||
);
|
||||
let error = result.unwrap_err();
|
||||
|
||||
// Check the error data contains the expected code
|
||||
match error {
|
||||
ServiceError::McpError(error_data) => {
|
||||
assert_eq!(
|
||||
error_data.code,
|
||||
ErrorCode::METHOD_NOT_FOUND,
|
||||
"Expected METHOD_NOT_FOUND error code (-32601)"
|
||||
);
|
||||
assert!(
|
||||
error_data
|
||||
.message
|
||||
.contains("requires task-based invocation"),
|
||||
"Error message should indicate task-based invocation is required, got: {}",
|
||||
error_data.message
|
||||
);
|
||||
}
|
||||
_ => panic!("Expected McpError variant, got: {:?}", error),
|
||||
}
|
||||
|
||||
client.cancel().await?;
|
||||
server_handle.await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_forbidden_task_tool_with_task_returns_error() -> anyhow::Result<()> {
|
||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||
|
||||
let server = TaskSupportTestServer::new();
|
||||
let server_handle = tokio::spawn(async move {
|
||||
server.serve(server_transport).await?.waiting().await?;
|
||||
anyhow::Ok(())
|
||||
});
|
||||
|
||||
let client_handler = DummyClientHandler::default();
|
||||
let client = client_handler.serve(client_transport).await?;
|
||||
|
||||
// Call the forbidden task tool WITH a task - should fail
|
||||
let result = client
|
||||
.call_tool(CallToolRequestParams {
|
||||
meta: None,
|
||||
name: "forbidden_task_tool".into(),
|
||||
arguments: None,
|
||||
task: make_task(), // Task provided but forbidden!
|
||||
})
|
||||
.await;
|
||||
|
||||
// Should be an error with code INVALID_PARAMS
|
||||
assert!(
|
||||
result.is_err(),
|
||||
"Expected error for forbidden task tool with task"
|
||||
);
|
||||
let error = result.unwrap_err();
|
||||
|
||||
// Check the error data contains the expected code
|
||||
match error {
|
||||
ServiceError::McpError(error_data) => {
|
||||
assert_eq!(
|
||||
error_data.code,
|
||||
ErrorCode::INVALID_PARAMS,
|
||||
"Expected INVALID_PARAMS error code"
|
||||
);
|
||||
assert!(
|
||||
error_data
|
||||
.message
|
||||
.contains("does not support task-based invocation"),
|
||||
"Error message should indicate task-based invocation is not supported, got: {}",
|
||||
error_data.message
|
||||
);
|
||||
}
|
||||
_ => panic!("Expected McpError variant, got: {:?}", error),
|
||||
}
|
||||
|
||||
client.cancel().await?;
|
||||
server_handle.await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_forbidden_task_tool_without_task_succeeds() -> anyhow::Result<()> {
|
||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||
|
||||
let server = TaskSupportTestServer::new();
|
||||
let server_handle = tokio::spawn(async move {
|
||||
server.serve(server_transport).await?.waiting().await?;
|
||||
anyhow::Ok(())
|
||||
});
|
||||
|
||||
let client_handler = DummyClientHandler::default();
|
||||
let client = client_handler.serve(client_transport).await?;
|
||||
|
||||
// Call the forbidden task tool WITHOUT a task - should succeed
|
||||
let result = client
|
||||
.call_tool(CallToolRequestParams {
|
||||
meta: None,
|
||||
name: "forbidden_task_tool".into(),
|
||||
arguments: None,
|
||||
task: None, // No task - allowed for forbidden
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"Forbidden task tool without task should succeed"
|
||||
);
|
||||
let result = result.unwrap();
|
||||
let text = result
|
||||
.content
|
||||
.first()
|
||||
.and_then(|c| c.raw.as_text())
|
||||
.map(|t| t.text.as_str())
|
||||
.unwrap_or("");
|
||||
assert_eq!(text, "forbidden task executed");
|
||||
|
||||
client.cancel().await?;
|
||||
server_handle.await??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_optional_task_tool_without_task_succeeds() -> anyhow::Result<()> {
|
||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||
|
||||
let server = TaskSupportTestServer::new();
|
||||
let server_handle = tokio::spawn(async move {
|
||||
server.serve(server_transport).await?.waiting().await?;
|
||||
anyhow::Ok(())
|
||||
});
|
||||
|
||||
let client_handler = DummyClientHandler::default();
|
||||
let client = client_handler.serve(client_transport).await?;
|
||||
|
||||
// Call the optional task tool WITHOUT a task - should succeed
|
||||
let result = client
|
||||
.call_tool(CallToolRequestParams {
|
||||
meta: None,
|
||||
name: "optional_task_tool".into(),
|
||||
arguments: None,
|
||||
task: None, // No task - allowed for optional
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"Optional task tool without task should succeed"
|
||||
);
|
||||
let result = result.unwrap();
|
||||
let text = result
|
||||
.content
|
||||
.first()
|
||||
.and_then(|c| c.raw.as_text())
|
||||
.map(|t| t.text.as_str())
|
||||
.unwrap_or("");
|
||||
assert_eq!(text, "optional task executed");
|
||||
|
||||
client.cancel().await?;
|
||||
server_handle.await??;
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -124,6 +124,7 @@ impl ServerHandler for SamplingDemoServer {
|
|||
),
|
||||
output_schema: None,
|
||||
annotations: None,
|
||||
execution: None,
|
||||
icons: None,
|
||||
meta: None,
|
||||
}],
|
||||
|
|
|
|||
Loading…
Reference in a new issue