modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat(task): add task support (SEP-1686) (#536)

Browse source 
Signed-off-by: jokemanfire <hu.dingyang@zte.com.cn>
This commit is contained in:
jokemanfire 2025-12-22 22:01:00 +08:00 • committed by GitHub
commit 621c9f619e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 1922 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

4
crates/rmcp/src/error.rs
View file

@ -41,6 +41,10 @@ pub enum RmcpError {
error: Box<dyn std::error::Error + Send + Sync>,
},
// and cancellation shouldn't be an error?
// TODO: add more error variants as needed
#[error("Task error: {0}")]
TaskError(String),
}
impl RmcpError {

77
crates/rmcp/src/handler/server.rs
View file

@ -11,6 +11,7 @@ pub mod router;
pub mod tool;
pub mod tool_name_validation;
pub mod wrapper;
impl<H: ServerHandler> Service<RoleServer> for H {
async fn handle_request(
&self,
@ -61,10 +62,18 @@ impl<H: ServerHandler> Service<RoleServer> for H {
.unsubscribe(request.params, context)
.await
.map(ServerResult::empty),
ClientRequest::CallToolRequest(request) => self
.call_tool(request.params, context)
.await
.map(ServerResult::CallToolResult),
ClientRequest::CallToolRequest(request) => {
if request.params.task.is_some() {
tracing::info!("Enqueueing task for tool call: {}", request.params.name);
self.enqueue_task(request.params, context.clone())
.await
.map(ServerResult::CreateTaskResult)
} else {
self.call_tool(request.params, context)
.await
.map(ServerResult::CallToolResult)
}
}
ClientRequest::ListToolsRequest(request) => self
.list_tools(request.params, context)
.await
@ -73,6 +82,22 @@ impl<H: ServerHandler> Service<RoleServer> for H {
.on_custom_request(request, context)
.await
.map(ServerResult::CustomResult),
ClientRequest::ListTasksRequest(request) => self
.list_tasks(request.params, context)
.await
.map(ServerResult::ListTasksResult),
ClientRequest::GetTaskInfoRequest(request) => self
.get_task_info(request.params, context)
.await
.map(ServerResult::GetTaskInfoResult),
ClientRequest::GetTaskResultRequest(request) => self
.get_task_result(request.params, context)
.await
.map(ServerResult::TaskResult),
ClientRequest::CancelTaskRequest(request) => self
.cancel_task(request.params, context)
.await
.map(ServerResult::empty),
}
}
@ -108,6 +133,16 @@ impl<H: ServerHandler> Service<RoleServer> for H {
#[allow(unused_variables)]
pub trait ServerHandler: Sized + Send + Sync + 'static {
fn enqueue_task(
&self,
_request: CallToolRequestParam,
_context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CreateTaskResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::internal_error(
"Task processing not implemented".to_string(),
None,
)))
}
fn ping(
&self,
context: RequestContext<RoleServer>,
@ -257,4 +292,38 @@ pub trait ServerHandler: Sized + Send + Sync + 'static {
fn get_info(&self) -> ServerInfo {
ServerInfo::default()
}
fn list_tasks(
&self,
request: Option<PaginatedRequestParam>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListTasksResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<ListTasksMethod>()))
}
fn get_task_info(
&self,
request: GetTaskInfoParam,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetTaskInfoResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<GetTaskInfoMethod>()))
}
fn get_task_result(
&self,
request: GetTaskResultParam,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<TaskResult, McpError>> + Send + '_ {
let _ = (request, context);
std::future::ready(Err(McpError::method_not_found::<GetTaskResultMethod>()))
}
fn cancel_task(
&self,
request: CancelTaskParam,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
let _ = (request, context);
std::future::ready(Err(McpError::method_not_found::<CancelTaskMethod>()))
}
}

8
crates/rmcp/src/handler/server/tool.rs
View file

@ -33,12 +33,17 @@ pub struct ToolCallContext<'s, S> {
pub service: &'s S,
pub name: Cow<'static, str>,
pub arguments: Option<JsonObject>,
pub task: Option<JsonObject>,
}
impl<'s, S> ToolCallContext<'s, S> {
pub fn new(
service: &'s S,
CallToolRequestParam { name, arguments }: CallToolRequestParam,
CallToolRequestParam {
name,
arguments,
task,
}: CallToolRequestParam,
request_context: RequestContext<RoleServer>,
) -> Self {
Self {
@ -46,6 +51,7 @@ impl<'s, S> ToolCallContext<'s, S> {
service,
name,
arguments,
task,
}
}
pub fn name(&self) -> &str {

2
crates/rmcp/src/lib.rs
View file

@ -125,6 +125,7 @@
//! .call_tool(CallToolRequestParam {
//! name: "git_status".into(),
//! arguments: serde_json::json!({ "repo_path": "." }).as_object().cloned(),
//! task: None,
//! })
//! .await?;
//! println!("Tool result: {tool_result:#?}");
@ -162,6 +163,7 @@ pub use service::{RoleClient, serve_client};
pub use service::{RoleServer, serve_server};
pub mod handler;
pub mod task_manager;
pub mod transport;
// re-export

73
crates/rmcp/src/model.rs
View file

@ -8,6 +8,7 @@ mod meta;
mod prompt;
mod resource;
mod serde_impl;
mod task;
mod tool;
pub use annotated::*;
pub use capabilities::*;
@ -19,6 +20,7 @@ pub use prompt::*;
pub use resource::*;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;
pub use task::*;
pub use tool::*;
/// A JSON object type alias for convenient handling of JSON data.
@ -1705,6 +1707,8 @@ pub struct CallToolRequestParam {
/// Arguments to pass to the tool (must match the tool's input schema)
#[serde(skip_serializing_if = "Option::is_none")]
pub arguments: Option<JsonObject>,
#[serde(skip_serializing_if = "Option::is_none")]
pub task: Option<JsonObject>,
}
/// Request to call a specific tool
@ -1743,6 +1747,61 @@ pub struct GetPromptResult {
pub messages: Vec<PromptMessage>,
}
// =============================================================================
// TASK MANAGEMENT
// =============================================================================
const_string!(GetTaskInfoMethod = "tasks/get");
pub type GetTaskInfoRequest = Request<GetTaskInfoMethod, GetTaskInfoParam>;
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct GetTaskInfoParam {
pub task_id: String,
}
const_string!(ListTasksMethod = "tasks/list");
pub type ListTasksRequest = RequestOptionalParam<ListTasksMethod, PaginatedRequestParam>;
const_string!(GetTaskResultMethod = "tasks/result");
pub type GetTaskResultRequest = Request<GetTaskResultMethod, GetTaskResultParam>;
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct GetTaskResultParam {
pub task_id: String,
}
const_string!(CancelTaskMethod = "tasks/cancel");
pub type CancelTaskRequest = Request<CancelTaskMethod, CancelTaskParam>;
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct CancelTaskParam {
pub task_id: String,
}
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct GetTaskInfoResult {
#[serde(skip_serializing_if = "Option::is_none")]
pub task: Option<crate::model::Task>,
}
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ListTasksResult {
pub tasks: Vec<crate::model::Task>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub total: Option<u64>,
}
// =============================================================================
// MESSAGE TYPE UNIONS
// =============================================================================
@ -1810,7 +1869,11 @@ ts_union!(
| UnsubscribeRequest
| CallToolRequest
| ListToolsRequest
| CustomRequest;
| CustomRequest
| GetTaskInfoRequest
| ListTasksRequest
| GetTaskResultRequest
| CancelTaskRequest;
);
impl ClientRequest {
@ -1830,6 +1893,10 @@ impl ClientRequest {
ClientRequest::CallToolRequest(r) => r.method.as_str(),
ClientRequest::ListToolsRequest(r) => r.method.as_str(),
ClientRequest::CustomRequest(r) => r.method.as_str(),
ClientRequest::GetTaskInfoRequest(r) => r.method.as_str(),
ClientRequest::ListTasksRequest(r) => r.method.as_str(),
ClientRequest::GetTaskResultRequest(r) => r.method.as_str(),
ClientRequest::CancelTaskRequest(r) => r.method.as_str(),
}
}
}
@ -1895,6 +1962,10 @@ ts_union!(
| CreateElicitationResult
| EmptyResult
| CustomResult
| CreateTaskResult
| ListTasksResult
| GetTaskInfoResult
| TaskResult
;
);

47
crates/rmcp/src/model/capabilities.rs
View file

@ -40,6 +40,25 @@ pub struct RootsCapabilities {
pub list_changed: Option<bool>,
}
/// Task capability negotiation for SEP-1686.
#[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`.
#[serde(skip_serializing_if = "Option::is_none")]
pub list: Option<bool>,
/// Whether the receiver supports `tasks/cancel`.
#[serde(skip_serializing_if = "Option::is_none")]
pub cancel: Option<bool>,
}
/// A convenience alias for describing per-request task support.
pub type TaskRequestMap = BTreeMap<String, bool>;
/// Capability for handling elicitation requests from servers.
///
/// Elicitation allows servers to request interactive input from users during tool execution.
@ -78,6 +97,8 @@ pub struct ClientCapabilities {
/// Capability to handle elicitation requests from servers for interactive user input
#[serde(skip_serializing_if = "Option::is_none")]
pub elicitation: Option<ElicitationCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tasks: Option<TasksCapability>,
}
///
@ -109,6 +130,8 @@ pub struct ServerCapabilities {
pub resources: Option<ResourcesCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<ToolsCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tasks: Option<TasksCapability>,
}
macro_rules! builder {
@ -223,12 +246,13 @@ builder! {
completions: JsonObject,
prompts: PromptsCapability,
resources: ResourcesCapability,
tools: ToolsCapability
tools: ToolsCapability,
tasks: TasksCapability
}
}
impl<const E: bool, const L: bool, const C: bool, const P: bool, const R: bool>
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, C, P, R, true>>
impl<const E: bool, const L: bool, const C: bool, const P: bool, const R: bool, const TASKS: bool>
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, C, P, R, true, TASKS>>
{
pub fn enable_tool_list_changed(mut self) -> Self {
if let Some(c) = self.tools.as_mut() {
@ -238,8 +262,8 @@ impl<const E: bool, const L: bool, const C: bool, const P: bool, const R: bool>
}
}
impl<const E: bool, const L: bool, const C: bool, const R: bool, const T: bool>
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, C, true, R, T>>
impl<const E: bool, const L: bool, const C: bool, const R: bool, const T: bool, const TASKS: bool>
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, C, true, R, T, TASKS>>
{
pub fn enable_prompts_list_changed(mut self) -> Self {
if let Some(c) = self.prompts.as_mut() {
@ -249,8 +273,8 @@ impl<const E: bool, const L: bool, const C: bool, const R: bool, const T: bool>
}
}
impl<const E: bool, const L: bool, const C: bool, const P: bool, const T: bool>
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, C, P, true, T>>
impl<const E: bool, const L: bool, const C: bool, const P: bool, const T: bool, const TASKS: bool>
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, C, P, true, T, TASKS>>
{
pub fn enable_resources_list_changed(mut self) -> Self {
if let Some(c) = self.resources.as_mut() {
@ -273,11 +297,12 @@ builder! {
roots: RootsCapabilities,
sampling: JsonObject,
elicitation: ElicitationCapability,
tasks: TasksCapability,
}
}
impl<const E: bool, const S: bool>
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, true, S>>
impl<const E: bool, const S: bool, const EL: bool, const TASKS: bool>
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, true, S, EL, TASKS>>
{
pub fn enable_roots_list_changed(mut self) -> Self {
if let Some(c) = self.roots.as_mut() {
@ -288,8 +313,8 @@ impl<const E: bool, const S: bool>
}
#[cfg(feature = "elicitation")]
impl<const E: bool, const R: bool, const S: bool>
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, R, S, true>>
impl<const E: bool, const R: bool, const S: bool, const TASKS: bool>
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, R, S, true, TASKS>>
{
/// Enable JSON Schema validation for elicitation responses.
/// When enabled, the client will validate user input against the requested_schema

4
crates/rmcp/src/model/meta.rs
View file

@ -107,6 +107,10 @@ variant_extension! {
CallToolRequest
ListToolsRequest
CustomRequest
GetTaskInfoRequest
ListTasksRequest
GetTaskResultRequest
CancelTaskRequest
}
}

79
crates/rmcp/src/model/task.rs Normal file
View file

@ -0,0 +1,79 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// Canonical task lifecycle status as defined by SEP-1686.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub enum TaskStatus {
/// The receiver accepted the request and is currently working on it.
#[default]
Working,
/// The receiver requires additional input before work can continue.
InputRequired,
/// The underlying operation completed successfully and the result is ready.
Completed,
/// The underlying operation failed and will not continue.
Failed,
/// The task was cancelled and will not continue processing.
Cancelled,
}
/// Final result for a succeeded task (returned from `tasks/result`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct TaskResult {
/// MIME type or custom content-type identifier.
pub content_type: String,
/// The actual result payload, matching the underlying request's schema.
pub value: Value,
/// Optional short summary for UI surfaces.
#[serde(skip_serializing_if = "Option::is_none")]
pub summary: Option<String>,
}
/// Primary Task object that surfaces metadata during the task lifecycle.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Task {
/// Unique task identifier generated by the receiver.
pub task_id: String,
/// Current lifecycle status (see [`TaskStatus`]).
pub status: TaskStatus,
/// Optional human-readable status message for UI surfaces.
#[serde(skip_serializing_if = "Option::is_none")]
pub status_message: Option<String>,
/// ISO-8601 creation timestamp.
pub created_at: String,
/// ISO-8601 timestamp for the most recent status change.
#[serde(skip_serializing_if = "Option::is_none")]
pub last_updated_at: Option<String>,
/// Retention window in milliseconds that the receiver agreed to honor.
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl: Option<u64>,
/// Suggested polling interval (milliseconds).
#[serde(skip_serializing_if = "Option::is_none")]
pub poll_interval: Option<u64>,
}
/// Wrapper returned by task-augmented requests (CreateTaskResult in SEP-1686).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct CreateTaskResult {
pub task: Task,
}
/// Paginated list of tasks
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct TaskList {
pub tasks: Vec<Task>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub total: Option<u64>,
}

294
crates/rmcp/src/task_manager.rs Normal file
View file

@ -0,0 +1,294 @@
use std::{any::Any, collections::HashMap, pin::Pin};
use futures::Future;
use tokio::{
sync::mpsc,
time::{Duration, timeout},
};
use crate::{
RoleServer,
error::{ErrorData as McpError, RmcpError as Error},
model::{CallToolResult, ClientRequest},
service::RequestContext,
};
/// Boxed future that represents an asynchronous operation managed by the processor.
pub type OperationFuture =
Pin<Box<dyn Future<Output = Result<Box<dyn OperationResultTransport>, Error>> + Send>>;
/// Describes metadata associated with an enqueued task.
#[derive(Debug, Clone)]
pub struct OperationDescriptor {
pub operation_id: String,
pub name: String,
pub client_request: Option<ClientRequest>,
pub context: Option<RequestContext<RoleServer>>,
pub ttl: Option<u64>,
}
impl OperationDescriptor {
pub fn new(operation_id: impl Into<String>, name: impl Into<String>) -> Self {
Self {
operation_id: operation_id.into(),
name: name.into(),
client_request: None,
context: None,
ttl: None,
}
}
pub fn with_client_request(mut self, request: ClientRequest) -> Self {
self.client_request = Some(request);
self
}
pub fn with_context(mut self, context: RequestContext<RoleServer>) -> Self {
self.context = Some(context);
self
}
pub fn with_ttl(mut self, ttl: u64) -> Self {
self.ttl = Some(ttl);
self
}
}
/// Operation message describing a unit of asynchronous work.
pub struct OperationMessage {
pub descriptor: OperationDescriptor,
pub future: OperationFuture,
}
impl OperationMessage {
pub fn new(descriptor: OperationDescriptor, future: OperationFuture) -> Self {
Self { descriptor, future }
}
}
/// Trait for operation result transport
pub trait OperationResultTransport: Send + Sync + 'static {
fn operation_id(&self) -> &String;
fn as_any(&self) -> &dyn std::any::Any;
}
// ===== Operation Processor =====
pub const DEFAULT_TASK_TIMEOUT_SECS: u64 = 300; // 5 minutes
/// Operation processor that coordinates extractors and handlers
pub struct OperationProcessor {
/// Currently running tasks keyed by id
running_tasks: HashMap<String, RunningTask>,
/// Completed results waiting to be collected
completed_results: Vec<TaskResult>,
task_result_receiver: Option<mpsc::UnboundedReceiver<TaskResult>>,
task_result_sender: mpsc::UnboundedSender<TaskResult>,
}
struct RunningTask {
task_handle: tokio::task::JoinHandle<()>,
started_at: std::time::Instant,
timeout: Option<u64>,
descriptor: OperationDescriptor,
}
pub struct TaskResult {
pub descriptor: OperationDescriptor,
pub result: Result<Box<dyn OperationResultTransport>, Error>,
}
/// Helper to generate an ISO 8601 timestamp for task metadata.
pub fn current_timestamp() -> String {
chrono::Utc::now().to_rfc3339()
}
/// Result transport for tool calls executed as tasks.
pub struct ToolCallTaskResult {
id: String,
pub result: Result<CallToolResult, McpError>,
}
impl ToolCallTaskResult {
pub fn new(id: impl Into<String>, result: Result<CallToolResult, McpError>) -> Self {
Self {
id: id.into(),
result,
}
}
}
impl OperationResultTransport for ToolCallTaskResult {
fn operation_id(&self) -> &String {
&self.id
}
fn as_any(&self) -> &dyn Any {
self
}
}
impl Default for OperationProcessor {
fn default() -> Self {
Self::new()
}
}
impl OperationProcessor {
pub fn new() -> Self {
let (task_result_sender, task_result_receiver) = mpsc::unbounded_channel();
Self {
running_tasks: HashMap::new(),
completed_results: Vec::new(),
task_result_receiver: Some(task_result_receiver),
task_result_sender,
}
}
/// Submit an operation for asynchronous execution.
#[allow(clippy::result_large_err)]
pub fn submit_operation(&mut self, message: OperationMessage) -> Result<(), Error> {
if self
.running_tasks
.contains_key(&message.descriptor.operation_id)
{
return Err(Error::TaskError(format!(
"Operation with id {} is already running",
message.descriptor.operation_id
)));
}
self.spawn_async_task(message);
Ok(())
}
fn spawn_async_task(&mut self, message: OperationMessage) {
let OperationMessage { descriptor, future } = message;
let task_id = descriptor.operation_id.clone();
let timeout_secs = descriptor.ttl.or(Some(DEFAULT_TASK_TIMEOUT_SECS));
let sender = self.task_result_sender.clone();
let descriptor_for_result = descriptor.clone();
let timed_future = async move {
if let Some(secs) = timeout_secs {
match timeout(Duration::from_secs(secs), future).await {
Ok(result) => result,
Err(_) => Err(Error::TaskError("Operation timed out".to_string())),
}
} else {
future.await
}
};
let handle = tokio::spawn(async move {
let result = timed_future.await;
let task_result = TaskResult {
descriptor: descriptor_for_result,
result,
};
let _ = sender.send(task_result);
});
let running_task = RunningTask {
task_handle: handle,
started_at: std::time::Instant::now(),
timeout: timeout_secs,
descriptor,
};
self.running_tasks.insert(task_id, running_task);
}
/// Collect completed results from running tasks and remove them from the running tasks map.
pub fn collect_completed_results(&mut self) -> Vec<TaskResult> {
if let Some(receiver) = &mut self.task_result_receiver {
while let Ok(result) = receiver.try_recv() {
self.running_tasks.remove(&result.descriptor.operation_id);
self.completed_results.push(result);
}
}
std::mem::take(&mut self.completed_results)
}
/// Check for tasks that have exceeded their timeout and handle them appropriately.
pub fn check_timeouts(&mut self) {
let now = std::time::Instant::now();
let mut timed_out_tasks = Vec::new();
for (task_id, task) in &self.running_tasks {
if let Some(timeout_duration) = task.timeout {
if now.duration_since(task.started_at).as_secs() > timeout_duration {
task.task_handle.abort();
timed_out_tasks.push(task_id.clone());
}
}
}
for task_id in timed_out_tasks {
if let Some(task) = self.running_tasks.remove(&task_id) {
let timeout_result = TaskResult {
descriptor: task.descriptor,
result: Err(Error::TaskError("Operation timed out".to_string())),
};
self.completed_results.push(timeout_result);
}
}
}
/// Get the number of running tasks.
pub fn running_task_count(&self) -> usize {
self.running_tasks.len()
}
/// Cancel all running tasks.
pub fn cancel_all_tasks(&mut self) {
for (_, task) in self.running_tasks.drain() {
task.task_handle.abort();
}
self.completed_results.clear();
}
/// List running task ids.
pub fn list_running(&self) -> Vec<String> {
self.running_tasks.keys().cloned().collect()
}
/// Note: collectors should call collect_completed_results; this provides a snapshot of queued results.
pub fn peek_completed(&self) -> &[TaskResult] {
&self.completed_results
}
/// Fetch the metadata for a running or recently completed task.
pub fn task_descriptor(&self, task_id: &str) -> Option<&OperationDescriptor> {
if let Some(task) = self.running_tasks.get(task_id) {
return Some(&task.descriptor);
}
self.completed_results
.iter()
.rev()
.find(|result| result.descriptor.operation_id == task_id)
.map(|result| &result.descriptor)
}
/// Attempt to cancel a running task.
pub fn cancel_task(&mut self, task_id: &str) -> bool {
if let Some(task) = self.running_tasks.remove(task_id) {
task.task_handle.abort();
// Insert a cancelled result so callers can observe the terminal state.
let cancel_result = TaskResult {
descriptor: task.descriptor,
result: Err(Error::TaskError("Operation cancelled".to_string())),
};
self.completed_results.push(cancel_result);
return true;
}
false
}
/// Retrieve a completed task result if available.
pub fn take_completed_result(&mut self, task_id: &str) -> Option<TaskResult> {
if let Some(position) = self
.completed_results
.iter()
.position(|result| result.descriptor.operation_id == task_id)
{
Some(self.completed_results.remove(position))
} else {
None
}
}
}

1
crates/rmcp/src/transport/streamable_http_client.rs
View file

@ -384,6 +384,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
"process initialized notification response",
))?;
let _ = initialized_notification.responder.send(Ok(()));
#[allow(clippy::large_enum_variant)]
enum Event<W: Worker, E: std::error::Error + Send + Sync + 'static> {
ClientMessage(WorkerSendRequest<W>),
ServerMessage(ServerJsonRpcMessage),