feat: Implement SEP-2663 Tasks Extension (#1020)
* feat!: implement SEP-2663 Tasks extension, removing the experimental 2025-11-25 tasks design Reshape tasks from the experimental SEP-1319/1686 core-protocol feature into the official io.modelcontextprotocol/tasks extension (SEP-2663): Model: - Re-model Task (statusMessage, ttlMs nullable, pollIntervalMs) and add DetailedTask with status-discriminated payloads (inputRequests/result/error inlined per spec) - CreateTaskResult now flattens Task with resultType: "task"; add ResultType::TASK and CallToolResponse::Task - Add tasks/update (UpdateTaskParams with MRTR InputResponses); tasks/get and tasks/cancel reworked; remove tasks/list, tasks/result - notifications/tasks now carries a full DetailedTask; removed from client notifications (SEP-2260) Capabilities: - Remove core TasksCapability et al; tasks are declared via the extensions map (enable_tasks() builders, supports_tasks() accessors) - Remove tool-level execution.taskSupport and the _meta.task / TaskMetadata / with_task() opt-in: task creation is server-directed and gated on the per-request client capability Runtime: - Replace OperationProcessor with TaskManager: durable-before-response task creation, input_required round-trips via tasks/update, cooperative cancellation, TTL expiry - Client peer helpers get_task/update_task/cancel_task - Emit Mcp-Name routing header from params.taskId for tasks/* methods (SEP-2243/2663) Macros: - Remove #[task_handler] and the tool execution() attribute Tests/examples/docs updated; schema snapshots regenerated. BREAKING CHANGE: the 2025-11-25 experimental tasks API is removed without a compatibility shim. Clients that do not declare the tasks extension always receive synchronous results. * feat: gate tasks/* methods on the client tasks-extension capability (SEP-2663) When the server advertises io.modelcontextprotocol/tasks but the client did not declare it (per-request _meta clientCapabilities, or initialize-time capabilities in session mode), tasks/get, tasks/update, and tasks/cancel now return -32021 Missing Required Client Capability with the required extension in the error data, instead of falling through to the handler's -32601 default. Servers that do not advertise the extension keep returning -32601. Also add regression tests confirming that unknown taskIds yield -32602 and that a legacy 2025-11-25 'task' param on tools/call is silently ignored. * feat: pass SEP-2663 Tasks extension conformance suite Conformance fixtures (conformance/src/bin/server.rs): - Add the required fixture tools: greet (sync-only), slow_compute, failing_job (task support: required), protocol_error_job, confirm_delete, multi_input, and test_tool_with_task (MRTR -> task escalation), all backed by TaskManager with server-directed task creation gated on the client's tasks-extension capability - Advertise io.modelcontextprotocol/tasks in server capabilities and wire get_task/update_task/cancel_task handlers - Task-required tools reject with -32021 when the client did not declare the extension SDK wire-shape fixes surfaced by the suite: - Add resultType: "complete" to GetTaskResult and introduce TaskAckResult so tasks/update and tasks/cancel acks carry the SEP-2322 discriminator (spec: every non-CreateTaskResult response on the tasks surface is resultType complete); dispatch now returns ServerResult::task_ack - CreateTaskResult gets a strict deserializer requiring resultType: "task" so it does not shadow task-shaped results in untagged unions - Client update_task/cancel_task accept both TaskAckResult and EmptyResult All 9 runnable Tasks extension server scenarios now pass (35/35 checks; tasks-status-notifications remains upstream-skipped), so the corresponding entries are removed from conformance/expected-failures-extensions.yaml. 2025-11-25 server suite (40/40), 2026-07-28 server suite (114/114), draft client suite, and extensions client suite all pass their baselines. * fix: address PR 1020 review feedback on DetailedTask schema and cooperative cancel - DetailedTask's JsonSchema now derives from the actual flattened wire shape (DetailedTaskWire: base Task + optional inputRequests/result/error) instead of approximating with the base Task schema, so generated schemas for GetTaskResult and notifications/tasks document the status-specific payload fields. Golden message schemas regenerated. - TaskManager::cancel_task no longer aborts the underlying future. It records the observable cancelled state and acks immediately (spec's eventually consistent semantics), but lets the operation keep running so it can observe is_cancel_requested() or the error from a woken request_input() and perform cleanup; any late result is discarded. New unit tests cover both the cooperative-cleanup path and waking parked input requests. * fix: make tasks/cancel truly cooperative per SEP-2663 (FEEDBACK_2) - TaskManager::cancel_task no longer forces terminal 'cancelled'. It records the cancellation intent, acks immediately, and wakes parked request_input awaits, but the operation decides its own terminal state: a post-cancel error settles as 'cancelled', while an operation that finishes its work settles as 'completed' — per the spec, 'the task may still reach a non-cancelled terminal status'. - Add TaskContext::cancelled(), a watch-based await for use with tokio::select! as the cooperative cancellation exit path. - Correct the unsupported-notification method string from 'notifications/tasks/status' to 'notifications/tasks' and document that task status notifications are not yet routable through subscriptions/listen (SubscriptionFilter has no taskIds field; upstream check still skipped). - Update conformance slow_compute fixture, task_demo example, and tests to honor cancellation via ctx.cancelled(); lifecycle conformance still 8/8 (all 9 scenarios remain green, 35/35 checks). * fix: sweep and evict expired tasks from TaskManager (FEEDBACK_3) TTL handling previously only ran from get_task and only flipped overdue non-terminal tasks to 'failed', never removing entries — an unbounded leak for long-lived servers, and it made ttl_ms: None = 'unlimited retention' meaningless since everything was retained forever. - Rename expire_overdue to sweep_expired and run it from every entry point (spawn, get_task, update_task, cancel_task). - Track terminal_at on each entry; terminal tasks are evicted after being retained for one further ttl_ms window past their terminal transition, so well-behaved pollers can observe the final state before late tasks/get calls return -32602 (spec: servers may delete expired tasks at any time, and returning task-not-found for purged tasks is compliant behavior). - ttl_ms: None entries are never evicted (spec: unlimited retention); document the retention model on TaskManager, including that there is no background sweeper. - New tests: retention-window eviction, sweep of abandoned tasks via other entry points, unlimited-TTL retention, and error-code assertions (-32602) for unknown ids across get/update/cancel. Reviewed the second feedback item (dedicated task-not-found error code) against SEP-2663 §Protocol Errors and it is incorrect: the SEP explicitly specifies -32602 (Invalid params) for invalid or nonexistent taskIds, which is what unknown_task already returns. Kept -32602; strengthened tests to assert the code. * fix: address Copilot review feedback on PR #1020 - TaskManager: drop the JoinHandle as soon as the operation settles instead of retaining it for the whole retention window, and only store it at spawn time if the task is still non-terminal (avoids keeping a completed handle that the completion path could never clear). - Use RequestContext::client_capabilities() (which applies the initialize-time fallback for session peers) instead of raw context.meta.client_capabilities() in the task_demo example, the README snippet, and the test server — the meta-only form would wrongly treat session-declared tasks clients as unsupported. All 9 Tasks extension conformance scenarios remain green (35/35 checks). * fix: enforce SEP-2663 task-result gating in tools/call dispatch A handler could return CallToolResponse::Task without checking whether the request declared the io.modelcontextprotocol/tasks extension, sending a CreateTaskResult to a client that cannot parse it. The SDK dispatch now rejects that case with -32021 Missing Required Client Capability before the response leaves the server, using RequestContext::client_capabilities() (per-request _meta with initialize-time fallback). Adds a regression test with a deliberately misbehaving handler that always materializes a task; all Tasks extension conformance scenarios remain green (35/35 checks). * fix: align TTL boundary comparisons in sweep_expired Copilot review: ttlMs: 0 never expired immediately because phase 1 used a strict 'elapsed > ttl_ms' comparison, and phase 2 retained terminal tasks at exactly the TTL boundary ('elapsed <= ttl_ms'). Treat elapsed >= ttl_ms as expired in phase 1 and evict when elapsed >= ttl_ms in phase 2, so both phases agree at the boundary and ttlMs: 0 expires/evicts on the first sweep. * fix: strict TaskAckResult deserializer and TASK_REQUIRED_TOOLS consistency - TaskAckResult carried only resultType (+ optional _meta) with a derived Deserialize, so inside the untagged ServerResult union it greedily matched any result object containing a resultType key, shadowing CustomResult and losing data (verified with a probe). Replace with a strict deserializer: deny_unknown_fields and require resultType == "complete". Regression tests cover both the non-matching shapes and the genuine ack shape. - Conformance fixtures: confirm_delete and multi_input rejected non-tasks clients inline, contradicting the TASK_REQUIRED_TOOLS doc/constant. Move them into TASK_REQUIRED_TOOLS (they park on in-task elicitation and have no synchronous fallback) so the upfront -32021 gate is the single source of truth, and drop the now-unreachable inline checks. All 9 Tasks extension conformance scenarios remain green (35/35 checks). * fix: let the operation decide its terminal state via TaskExit After tasks/cancel, any operation error was coerced to terminal 'cancelled', masking real failures (e.g. an unrelated error landing just after a late cancel request) and contradicting the documented 'operation decides its own terminal state' contract. Change TaskFuture's error type from McpError to a new TaskExit enum: - TaskExit::Cancelled — an explicit cooperative-cancellation exit; settles as terminal 'cancelled'. - TaskExit::Error(McpError) — a real failure; settles as terminal 'failed' with the error inlined, even after tasks/cancel was received. From<McpError> for TaskExit keeps '?' ergonomic in task bodies, and request_input() now returns TaskExit directly (its wake-on-cancel path yields TaskExit::Cancelled), so parked operations that propagate it with '?' settle as 'cancelled' automatically. Update the conformance fixtures, task_demo example, and tests; add a regression test asserting a post-cancel unrelated error settles as 'failed' with its error payload preserved. All 9 Tasks extension conformance scenarios remain green (35/35 checks). * fix: allow TaskExit enum to be exhaustive * fix: close spawn/shutdown race, sweep in running_task_count, clarify TTL retention docs Addresses branch-review findings: - spawn() raced with shutdown(): if shutdown drained the task map between the entry insert and the JoinHandle store, the handle was dropped and the operation kept running detached. If the entry is gone at store time, the handle is now aborted instead. - running_task_count() now runs the TTL sweep like every other entry point, so it no longer reports overdue tasks as running (matching the documented sweep-on-every-entry-point behavior). Regression test added. - Document that terminal-task retention intentionally extends one ttl_ms window past the terminal transition (observation grace period) beyond the creation-based lifetime ttlMs advertises on the wire — compliant since SEP-2663 allows deleting expired tasks at any time after the TTL.
This commit is contained in:
Jack Amadeo
GitHub
parent
ac9c637d04
commit
07abfdfd99
5 changed files with 2950 additions and 3236 deletions
|
|
@ -69,12 +69,13 @@ A client demonstrating how to use the sampling tool.
|
|||
|
||||
### Task Standard I/O Client (`task_stdio.rs`)
|
||||
|
||||
A client that exercises the task lifecycle against `servers_task_stdio`
|
||||
(per [SEP-1319](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks)).
|
||||
A client that exercises the SEP-2663 Tasks extension lifecycle against `servers_task_stdio`
|
||||
([SEP-2663](https://modelcontextprotocol.io/extensions/tasks/overview), `io.modelcontextprotocol/tasks`).
|
||||
|
||||
- Spawns `servers_task_stdio` as a child process over stdio
|
||||
- Declares the tasks extension in its client capabilities
|
||||
- Calls `quick_echo` synchronously
|
||||
- Calls `slow_sum` as a task via `CallToolRequestParams::with_task(...)`, polls `tasks/get` until completion, then fetches the result via `tasks/result`
|
||||
- Calls `slow_sum`, receives a `CreateTaskResult` (`resultType: "task"`), polls `tasks/get` honoring `pollIntervalMs`, and reads the final `CallToolResult` inlined in the completed task
|
||||
|
||||
### Progress Test Client (`progress_client.rs`)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,18 +1,20 @@
|
|||
//! Client for the task-demo server in `examples/servers/src/task_stdio.rs`.
|
||||
//!
|
||||
//! Walks through the task lifecycle (SEP-1319):
|
||||
//! Walks through the SEP-2663 Tasks extension lifecycle:
|
||||
//! 1. Call a regular tool (`quick_echo`) — synchronous response.
|
||||
//! 2. Call a task-required tool (`slow_sum`) by attaching `task: {}` to
|
||||
//! the `tools/call` request. The server returns a `Task` with a `task_id`.
|
||||
//! 3. Poll `tasks/get` until status becomes `Completed`.
|
||||
//! 4. Fetch the underlying `CallToolResult` via `tasks/result`.
|
||||
//! 2. Call `slow_sum` while declaring the `io.modelcontextprotocol/tasks`
|
||||
//! extension capability. The server decides to materialize a task and
|
||||
//! returns a `CreateTaskResult` (`resultType: "task"`).
|
||||
//! 3. Poll `tasks/get` (honoring `pollIntervalMs`) until the task reaches a
|
||||
//! terminal status; the final `CallToolResult` is inlined in the
|
||||
//! `completed` task's `result` field.
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use rmcp::{
|
||||
ServiceExt,
|
||||
model::{
|
||||
CallToolRequestParams, CallToolResult, ClientRequest, GetTaskParams, GetTaskPayloadParams,
|
||||
Request, ServerResult, TaskMetadata, TaskStatus,
|
||||
CallToolRequestParams, CallToolResponse, CallToolResult, ClientCapabilities, GetTaskParams,
|
||||
TaskPayload, TaskStatus,
|
||||
},
|
||||
object,
|
||||
transport::{ConfigureCommandExt, TokioChildProcess},
|
||||
|
|
@ -30,8 +32,14 @@ async fn main() -> Result<()> {
|
|||
.with(tracing_subscriber::fmt::layer())
|
||||
.init();
|
||||
|
||||
// Declare the tasks extension in our client capabilities (SEP-2663).
|
||||
let client_info = rmcp::model::ClientInfo::new(
|
||||
ClientCapabilities::builder().enable_tasks().build(),
|
||||
rmcp::model::Implementation::from_build_env(),
|
||||
);
|
||||
|
||||
// Spawn the task-demo server as a child process over stdio.
|
||||
let client = ()
|
||||
let client = client_info
|
||||
.serve(TokioChildProcess::new(Command::new("cargo").configure(
|
||||
|cmd| {
|
||||
cmd.arg("run")
|
||||
|
|
@ -44,7 +52,7 @@ async fn main() -> Result<()> {
|
|||
))?)
|
||||
.await?;
|
||||
|
||||
// 1) Synchronous call. `quick_echo` has the default task_support = forbidden.
|
||||
// 1) Synchronous call.
|
||||
let echo = client
|
||||
.call_tool(
|
||||
CallToolRequestParams::new("quick_echo")
|
||||
|
|
@ -53,68 +61,63 @@ async fn main() -> Result<()> {
|
|||
.await?;
|
||||
tracing::info!("quick_echo -> {echo:#?}");
|
||||
|
||||
// 2) Task call. `slow_sum` is task_support = required, so we MUST attach
|
||||
// `task` metadata. An empty `TaskMetadata` is fine; use `.with_ttl(...)`
|
||||
// to set a retention window.
|
||||
let create = client
|
||||
.send_request(ClientRequest::CallToolRequest(Request::new(
|
||||
CallToolRequestParams::new("slow_sum")
|
||||
.with_arguments(object!({ "a": 40, "b": 2 }))
|
||||
.with_task(TaskMetadata::new()),
|
||||
)))
|
||||
// 2) Task-eligible call. The server sees our tasks capability and
|
||||
// materializes a task instead of blocking.
|
||||
let response = client
|
||||
.call_tool_once(
|
||||
CallToolRequestParams::new("slow_sum").with_arguments(object!({ "a": 40, "b": 2 })),
|
||||
)
|
||||
.await?;
|
||||
let ServerResult::CreateTaskResult(create) = create else {
|
||||
return Err(anyhow!("expected CreateTaskResult, got {create:?}"));
|
||||
let create = match response {
|
||||
CallToolResponse::Task(create) => create,
|
||||
CallToolResponse::Complete(result) => {
|
||||
// The server is allowed to answer synchronously.
|
||||
tracing::info!("slow_sum answered synchronously -> {result:#?}");
|
||||
client.cancel().await?;
|
||||
return Ok(());
|
||||
}
|
||||
other => return Err(anyhow!("unexpected response: {other:?}")),
|
||||
};
|
||||
let task_id = create.task.task_id.clone();
|
||||
let poll_ms = create.task.poll_interval_ms.unwrap_or(500);
|
||||
tracing::info!(
|
||||
"slow_sum enqueued as task {task_id} (status = {:?})",
|
||||
"slow_sum materialized as task {task_id} (status = {:?})",
|
||||
create.task.status
|
||||
);
|
||||
|
||||
// 3) Poll `tasks/get` until the server reports a terminal status.
|
||||
let final_status = loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
||||
// 3) Poll `tasks/get` until the task reaches a terminal status.
|
||||
let final_task = loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(poll_ms)).await;
|
||||
|
||||
let info = client
|
||||
.send_request(ClientRequest::GetTaskRequest(Request::new(
|
||||
GetTaskParams::new(task_id.clone()),
|
||||
)))
|
||||
.peer()
|
||||
.get_task(GetTaskParams::new(task_id.clone()))
|
||||
.await?;
|
||||
let ServerResult::GetTaskResult(info) = info else {
|
||||
return Err(anyhow!("expected GetTaskResult, got {info:?}"));
|
||||
};
|
||||
tracing::info!("status = {:?}", info.task.status);
|
||||
tracing::info!("status = {:?}", info.task.status());
|
||||
|
||||
match info.task.status {
|
||||
TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled => {
|
||||
break info.task.status;
|
||||
}
|
||||
_ => {}
|
||||
if info.task.status().is_terminal() {
|
||||
break info.task;
|
||||
}
|
||||
};
|
||||
|
||||
if final_status != TaskStatus::Completed {
|
||||
return Err(anyhow!("task ended in {final_status:?}"));
|
||||
// The completed task carries the final CallToolResult inline.
|
||||
match &final_task.payload {
|
||||
TaskPayload::Completed { result } => {
|
||||
let call_result: CallToolResult =
|
||||
serde_json::from_value(serde_json::Value::Object(result.clone()))?;
|
||||
tracing::info!("slow_sum result -> {call_result:#?}");
|
||||
}
|
||||
TaskPayload::Failed { error } => {
|
||||
return Err(anyhow!("task failed: {error:?}"));
|
||||
}
|
||||
other => {
|
||||
return Err(anyhow!(
|
||||
"task ended in unexpected state {:?}",
|
||||
other.status()
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// 4) Fetch the payload. The server-side handler returns a serialized
|
||||
// `CallToolResult`. On the wire the response is just a JSON value, and
|
||||
// `ServerResult` is `#[serde(untagged)]`, so the client decodes it as
|
||||
// whichever variant the JSON shape matches first — a `CallToolResult`
|
||||
// here. (For a non-tool task the same value would surface as
|
||||
// `ServerResult::CustomResult` and need manual `serde_json::from_value`.)
|
||||
let payload = client
|
||||
.send_request(ClientRequest::GetTaskPayloadRequest(Request::new(
|
||||
GetTaskPayloadParams::new(task_id.clone()),
|
||||
)))
|
||||
.await?;
|
||||
let call_result: CallToolResult = match payload {
|
||||
ServerResult::CallToolResult(r) => r,
|
||||
ServerResult::CustomResult(c) => serde_json::from_value(c.0)?,
|
||||
other => return Err(anyhow!("unexpected task result: {other:?}")),
|
||||
};
|
||||
tracing::info!("slow_sum result -> {call_result:#?}");
|
||||
debug_assert_eq!(final_task.status(), TaskStatus::Completed);
|
||||
|
||||
client.cancel().await?;
|
||||
Ok(())
|
||||
|
|
|
|||
|
|
@ -71,13 +71,13 @@ A server demonstrating the prompt framework capabilities.
|
|||
|
||||
### Task Demo Server (`task_stdio.rs`)
|
||||
|
||||
A minimal stdio server demonstrating task-based tool invocation per
|
||||
[SEP-1319](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks).
|
||||
A minimal stdio server demonstrating the MCP Tasks extension
|
||||
([SEP-2663](https://modelcontextprotocol.io/extensions/tasks/overview), `io.modelcontextprotocol/tasks`).
|
||||
|
||||
- `slow_sum` is declared with `execution(task_support = "required")`, so clients MUST invoke it as a task
|
||||
- `slow_sum` is materialized as a task (`CreateTaskResult`, `resultType: "task"`) whenever the client declares the tasks extension capability; other clients get a normal synchronous response
|
||||
- `quick_echo` is a regular synchronous tool for contrast
|
||||
- Wires up `enqueue_task` / `tasks/get` / `tasks/result` / `tasks/cancel` via `#[task_handler]`
|
||||
- Pair with `examples/clients/src/task_stdio.rs` to see the full lifecycle (create → poll → fetch result)
|
||||
- Serves `tasks/get` / `tasks/update` / `tasks/cancel` via a `TaskManager`
|
||||
- Pair with `examples/clients/src/task_stdio.rs` to see the full lifecycle (create → poll → inline result)
|
||||
|
||||
### MRTR Demo (`mrtr.rs`)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
#![allow(dead_code)]
|
||||
use std::{any::Any, sync::Arc};
|
||||
use std::sync::Arc;
|
||||
|
||||
use rmcp::{
|
||||
ErrorData as McpError, RoleServer, ServerHandler,
|
||||
|
|
@ -10,28 +10,11 @@ use rmcp::{
|
|||
model::*,
|
||||
prompt, prompt_handler, prompt_router, schemars,
|
||||
service::RequestContext,
|
||||
task_handler,
|
||||
task_manager::{OperationProcessor, OperationResultTransport},
|
||||
tool, tool_handler, tool_router,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
struct ToolCallOperationResult {
|
||||
id: String,
|
||||
result: Result<CallToolResult, McpError>,
|
||||
}
|
||||
|
||||
impl OperationResultTransport for ToolCallOperationResult {
|
||||
fn operation_id(&self) -> &String {
|
||||
&self.id
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn Any {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct StructRequest {
|
||||
pub a: i32,
|
||||
|
|
@ -78,7 +61,6 @@ pub struct Counter {
|
|||
counter: Arc<Mutex<i32>>,
|
||||
tool_router: ToolRouter<Counter>,
|
||||
prompt_router: PromptRouter<Counter>,
|
||||
processor: Arc<Mutex<OperationProcessor>>,
|
||||
}
|
||||
|
||||
#[tool_router]
|
||||
|
|
@ -89,7 +71,6 @@ impl Counter {
|
|||
counter: Arc::new(Mutex::new(0)),
|
||||
tool_router: Self::tool_router(),
|
||||
prompt_router: Self::prompt_router(),
|
||||
processor: Arc::new(Mutex::new(OperationProcessor::new())),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -123,10 +104,7 @@ impl Counter {
|
|||
)]))
|
||||
}
|
||||
|
||||
#[tool(
|
||||
description = "Long running task example",
|
||||
execution(task_support = "optional")
|
||||
)]
|
||||
#[tool(description = "Long running task example")]
|
||||
async fn long_task(&self) -> Result<CallToolResult, McpError> {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(10)).await;
|
||||
Ok(CallToolResult::success(vec![ContentBlock::text(
|
||||
|
|
@ -227,7 +205,6 @@ impl Counter {
|
|||
|
||||
#[tool_handler(meta = MetaObject(rmcp::object!({"tool_meta_key": "tool_meta_value"})))]
|
||||
#[prompt_handler(meta = MetaObject(rmcp::object!({"router_meta_key": "router_meta_value"})))]
|
||||
#[task_handler]
|
||||
impl ServerHandler for Counter {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo::new(
|
||||
|
|
@ -380,49 +357,4 @@ mod tests {
|
|||
let prompts = router.list_all();
|
||||
assert_eq!(prompts.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_client_enqueues_long_task() -> anyhow::Result<()> {
|
||||
let counter = Counter::new();
|
||||
let processor = counter.processor.clone();
|
||||
let client = TestClient::default();
|
||||
|
||||
let (server_transport, client_transport) = tokio::io::duplex(4096);
|
||||
let server_handle = tokio::spawn(async move {
|
||||
let service = counter.serve(server_transport).await?;
|
||||
service.waiting().await?;
|
||||
anyhow::Ok(())
|
||||
});
|
||||
|
||||
let client_service = client.serve(client_transport).await?;
|
||||
let params = CallToolRequestParams::new("long_task").with_task(TaskMetadata::new());
|
||||
let response = client_service
|
||||
.send_request(ClientRequest::CallToolRequest(Request::new(params.clone())))
|
||||
.await?;
|
||||
|
||||
let ServerResult::CreateTaskResult(info) = response else {
|
||||
panic!("expected task creation result, got {response:?}");
|
||||
};
|
||||
let task = info.task;
|
||||
|
||||
assert_eq!(task.status, TaskStatus::Working);
|
||||
// task list should show the task
|
||||
let tasks = client_service
|
||||
.send_request(ClientRequest::ListTasksRequest(
|
||||
RequestOptionalParam::default(),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let ServerResult::ListTasksResult(listed) = tasks else {
|
||||
panic!("expected list tasks result, got {tasks:?}");
|
||||
};
|
||||
assert_eq!(listed.tasks[0].task_id, task.task_id);
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
let running = processor.lock().await.running_task_count();
|
||||
assert_eq!(running, 1);
|
||||
|
||||
client_service.cancel().await?;
|
||||
let _ = server_handle.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,27 +1,26 @@
|
|||
//! Minimal example of a tool that supports task-based invocation (SEP-1319).
|
||||
//! Minimal example of a server that supports the MCP Tasks extension
|
||||
//! (SEP-2663, `io.modelcontextprotocol/tasks`).
|
||||
//!
|
||||
//! - `slow_sum` is marked `task_support = "required"`, so the client MUST invoke
|
||||
//! it as a task. The server enqueues the call into an `OperationProcessor`,
|
||||
//! returns a task id immediately, and the client polls `tasks/get` and
|
||||
//! fetches the payload via `tasks/result`.
|
||||
//! - `quick_echo` is a regular synchronous tool for contrast (the default,
|
||||
//! `task_support = "forbidden"`).
|
||||
//! - `slow_sum` is executed as a task whenever the client declares the tasks
|
||||
//! extension capability: the server returns a `CreateTaskResult`
|
||||
//! (`resultType: "task"`) immediately and the client polls `tasks/get`.
|
||||
//! Clients that do not declare the extension get a normal synchronous
|
||||
//! response.
|
||||
//! - `quick_echo` is a regular synchronous tool for contrast.
|
||||
//!
|
||||
//! See `examples/clients/src/task_stdio.rs` for the matching client.
|
||||
|
||||
#![allow(dead_code)]
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use rmcp::{
|
||||
ErrorData as McpError, ServerHandler,
|
||||
handler::server::{router::tool::ToolRouter, wrapper::Parameters},
|
||||
model::{CallToolResult, ContentBlock},
|
||||
schemars, task_handler,
|
||||
task_manager::OperationProcessor,
|
||||
tool, tool_handler, tool_router,
|
||||
model::*,
|
||||
schemars,
|
||||
service::{RequestContext, RoleServer},
|
||||
task_manager::{TaskExit, TaskManager, TaskOptions},
|
||||
tool, tool_router,
|
||||
};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct SumArgs {
|
||||
|
|
@ -34,13 +33,10 @@ pub struct EchoArgs {
|
|||
pub message: String,
|
||||
}
|
||||
|
||||
/// Server state. The `processor` field is required by `#[task_handler]`:
|
||||
/// the macro generates `enqueue_task` / `tasks/*` handlers that submit and
|
||||
/// poll operations through it.
|
||||
#[derive(Clone)]
|
||||
pub struct TaskDemo {
|
||||
tool_router: ToolRouter<TaskDemo>,
|
||||
processor: Arc<Mutex<OperationProcessor>>,
|
||||
tasks: TaskManager,
|
||||
}
|
||||
|
||||
impl Default for TaskDemo {
|
||||
|
|
@ -54,17 +50,12 @@ impl TaskDemo {
|
|||
pub fn new() -> Self {
|
||||
Self {
|
||||
tool_router: Self::tool_router(),
|
||||
processor: Arc::new(Mutex::new(OperationProcessor::new())),
|
||||
tasks: TaskManager::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Long-running tool. The `execution(task_support = "required")` attribute
|
||||
/// tells clients they MUST call this tool as a task; the server returns
|
||||
/// `-32601` if they don't.
|
||||
#[tool(
|
||||
description = "Sum two numbers after a 2-second delay",
|
||||
execution(task_support = "required")
|
||||
)]
|
||||
/// Long-running tool. Run as a task when the client supports tasks.
|
||||
#[tool(description = "Sum two numbers after a 2-second delay")]
|
||||
async fn slow_sum(
|
||||
&self,
|
||||
Parameters(SumArgs { a, b }): Parameters<SumArgs>,
|
||||
|
|
@ -75,7 +66,7 @@ impl TaskDemo {
|
|||
)]))
|
||||
}
|
||||
|
||||
/// Synchronous tool with the default `task_support = "forbidden"`.
|
||||
/// Synchronous tool.
|
||||
#[tool(description = "Echo a message back immediately")]
|
||||
async fn quick_echo(
|
||||
&self,
|
||||
|
|
@ -85,9 +76,92 @@ impl TaskDemo {
|
|||
}
|
||||
}
|
||||
|
||||
/// `#[task_handler]` reads `self.processor` (configurable via the macro's
|
||||
/// `processor = ...` argument) and synthesizes `enqueue_task`, `list_tasks`,
|
||||
/// `get_task_info`, `get_task_result`, and `cancel_task` for us.
|
||||
#[tool_handler]
|
||||
#[task_handler]
|
||||
impl ServerHandler for TaskDemo {}
|
||||
impl ServerHandler for TaskDemo {
|
||||
async fn call_tool(
|
||||
&self,
|
||||
request: CallToolRequestParams,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> Result<CallToolResponse, McpError> {
|
||||
// SEP-2663: the server decides per-request whether to materialize a
|
||||
// task, but MUST NOT return one unless the request declared the tasks
|
||||
// extension capability.
|
||||
let client_supports_tasks = context
|
||||
.client_capabilities()
|
||||
.is_some_and(|caps| caps.supports_tasks());
|
||||
|
||||
if request.name == "slow_sum" && client_supports_tasks {
|
||||
let params: SumArgs = serde_json::from_value(serde_json::Value::Object(
|
||||
request.arguments.clone().unwrap_or_default(),
|
||||
))
|
||||
.map_err(|e| McpError::invalid_params(e.to_string(), None))?;
|
||||
let task = self.tasks.spawn(TaskOptions::default(), move |ctx| {
|
||||
Box::pin(async move {
|
||||
// Cancellation is cooperative (SEP-2663): honor
|
||||
// tasks/cancel by exiting with TaskExit::Cancelled.
|
||||
tokio::select! {
|
||||
_ = ctx.cancelled() => {
|
||||
Err(TaskExit::Cancelled)
|
||||
}
|
||||
_ = tokio::time::sleep(std::time::Duration::from_secs(2)) => {
|
||||
Ok(CallToolResult::success(vec![ContentBlock::text(
|
||||
(params.a + params.b).to_string(),
|
||||
)]))
|
||||
}
|
||||
}
|
||||
})
|
||||
});
|
||||
return Ok(CallToolResponse::Task(CreateTaskResult::new(task)));
|
||||
}
|
||||
|
||||
// Fall back to synchronous execution via the tool router.
|
||||
let tcc = rmcp::handler::server::tool::ToolCallContext::new(self, request, context);
|
||||
self.tool_router.call(tcc).await
|
||||
}
|
||||
|
||||
async fn list_tools(
|
||||
&self,
|
||||
_request: Option<PaginatedRequestParams>,
|
||||
_context: RequestContext<RoleServer>,
|
||||
) -> Result<ListToolsResult, McpError> {
|
||||
Ok(ListToolsResult::with_all_items(self.tool_router.list_all()))
|
||||
}
|
||||
|
||||
async fn get_task(
|
||||
&self,
|
||||
request: GetTaskParams,
|
||||
_context: RequestContext<RoleServer>,
|
||||
) -> Result<GetTaskResult, McpError> {
|
||||
Ok(GetTaskResult::new(self.tasks.get_task(&request.task_id)?))
|
||||
}
|
||||
|
||||
async fn update_task(
|
||||
&self,
|
||||
request: UpdateTaskParams,
|
||||
_context: RequestContext<RoleServer>,
|
||||
) -> Result<(), McpError> {
|
||||
self.tasks
|
||||
.update_task(&request.task_id, request.input_responses)
|
||||
}
|
||||
|
||||
async fn cancel_task(
|
||||
&self,
|
||||
request: CancelTaskParams,
|
||||
_context: RequestContext<RoleServer>,
|
||||
) -> Result<(), McpError> {
|
||||
self.tasks.cancel_task(&request.task_id)
|
||||
}
|
||||
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo::new(
|
||||
ServerCapabilities::builder()
|
||||
.enable_tools()
|
||||
.enable_tasks()
|
||||
.build(),
|
||||
)
|
||||
.with_instructions(
|
||||
"Task demo server (SEP-2663). `slow_sum` runs as a task for \
|
||||
clients that declare the tasks extension."
|
||||
.to_string(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue