chore(example): add example in progress (#271)
Signed-off-by: jokemanfire <hu.dingyang@zte.com.cn>
This commit is contained in:
parent
271cf0b232
commit
f8a747eaca
8 changed files with 726 additions and 2 deletions
|
|
@ -28,6 +28,7 @@ url = "2.4"
|
||||||
tower = "0.5"
|
tower = "0.5"
|
||||||
axum = "0.8"
|
axum = "0.8"
|
||||||
reqwest = "0.12"
|
reqwest = "0.12"
|
||||||
|
clap = { version = "4.0", features = ["derive"] }
|
||||||
|
|
||||||
[[example]]
|
[[example]]
|
||||||
name = "clients_sse"
|
name = "clients_sse"
|
||||||
|
|
@ -56,3 +57,8 @@ path = "src/auth/oauth_client.rs"
|
||||||
[[example]]
|
[[example]]
|
||||||
name = "clients_sampling_stdio"
|
name = "clients_sampling_stdio"
|
||||||
path = "src/sampling_stdio.rs"
|
path = "src/sampling_stdio.rs"
|
||||||
|
|
||||||
|
|
||||||
|
[[example]]
|
||||||
|
name = "clients_progress_client"
|
||||||
|
path = "src/progress_client.rs"
|
||||||
|
|
@ -4,6 +4,7 @@ This directory contains Model Context Protocol (MCP) client examples implemented
|
||||||
|
|
||||||
## Example List
|
## Example List
|
||||||
|
|
||||||
|
|
||||||
### SSE Client (`sse.rs`)
|
### SSE Client (`sse.rs`)
|
||||||
|
|
||||||
A client that communicates with an MCP server using Server-Sent Events (SSE) transport.
|
A client that communicates with an MCP server using Server-Sent Events (SSE) transport.
|
||||||
|
|
@ -67,6 +68,16 @@ A client demonstrating how to use the sampling tool.
|
||||||
- Retrieves server information and list of available tools
|
- Retrieves server information and list of available tools
|
||||||
- Calls the `ask_llm` tool
|
- Calls the `ask_llm` tool
|
||||||
|
|
||||||
|
### Progress Test Client (`progress_test_client.rs`)
|
||||||
|
|
||||||
|
A client that communicates with an MCP server using progress notifications.
|
||||||
|
|
||||||
|
- Launches the `cargo run --example clients_progress_client -- --transport {stdio|sse|http|all}` to test the progress notifications
|
||||||
|
- Connects to the server using different transport methods
|
||||||
|
- Tests the progress notifications
|
||||||
|
- The sse and http should run the server first
|
||||||
|
|
||||||
|
|
||||||
## How to Run
|
## How to Run
|
||||||
|
|
||||||
Each example can be run using Cargo:
|
Each example can be run using Cargo:
|
||||||
|
|
|
||||||
397
examples/clients/src/progress_client.rs
Normal file
397
examples/clients/src/progress_client.rs
Normal file
|
|
@ -0,0 +1,397 @@
|
||||||
|
use std::{
|
||||||
|
sync::{Arc, Mutex},
|
||||||
|
time::{Duration, Instant},
|
||||||
|
};
|
||||||
|
|
||||||
|
use anyhow::Result;
|
||||||
|
use clap::{Parser, ValueEnum};
|
||||||
|
use rmcp::{
|
||||||
|
ClientHandler, ServiceExt,
|
||||||
|
model::{
|
||||||
|
CallToolRequestParam, ClientCapabilities, ClientInfo, Implementation,
|
||||||
|
ProgressNotificationParam,
|
||||||
|
},
|
||||||
|
service::{NotificationContext, RoleClient},
|
||||||
|
transport::{SseClientTransport, StreamableHttpClientTransport, TokioChildProcess},
|
||||||
|
};
|
||||||
|
use tokio::{process::Command, time::sleep};
|
||||||
|
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, ValueEnum)]
|
||||||
|
enum TransportType {
|
||||||
|
Stdio,
|
||||||
|
Sse,
|
||||||
|
Http,
|
||||||
|
All,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Parser)]
|
||||||
|
#[command(name = "progress-test-client")]
|
||||||
|
#[command(about = "RMCP Progress Notification Test Client")]
|
||||||
|
struct Args {
|
||||||
|
/// Transport type to test
|
||||||
|
#[arg(short, long, value_enum, default_value_t = TransportType::Stdio)]
|
||||||
|
transport: TransportType,
|
||||||
|
|
||||||
|
/// Number of records to process
|
||||||
|
#[arg(short, long, default_value_t = 10)]
|
||||||
|
records: u32,
|
||||||
|
|
||||||
|
/// SSE server URL
|
||||||
|
#[arg(long, default_value = "http://127.0.0.1:8000/sse")]
|
||||||
|
sse_url: String,
|
||||||
|
|
||||||
|
/// HTTP server URL
|
||||||
|
#[arg(long, default_value = "http://127.0.0.1:8001/mcp")]
|
||||||
|
http_url: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct ProgressTracker {
|
||||||
|
start_time: Instant,
|
||||||
|
progress_count: Arc<Mutex<u32>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProgressTracker {
|
||||||
|
fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
start_time: Instant::now(),
|
||||||
|
progress_count: Arc::new(Mutex::new(0)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn handle_progress(&self, params: &ProgressNotificationParam) {
|
||||||
|
if let Ok(mut count) = self.progress_count.lock() {
|
||||||
|
*count += 1;
|
||||||
|
}
|
||||||
|
let elapsed = self.start_time.elapsed();
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"Progress update [{}]: {}/{} - {} (elapsed: {:.1}s)",
|
||||||
|
params.progress_token.0,
|
||||||
|
params.progress,
|
||||||
|
params.total.unwrap_or(0.0),
|
||||||
|
params.message.as_deref().unwrap_or(""),
|
||||||
|
elapsed.as_secs_f64()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn print_summary(&self) {
|
||||||
|
let elapsed = self.start_time.elapsed();
|
||||||
|
if let Ok(count) = self.progress_count.lock() {
|
||||||
|
tracing::info!("Total progress notifications received: {}", *count);
|
||||||
|
}
|
||||||
|
tracing::info!("Total time elapsed: {:.2}s", elapsed.as_secs_f64());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Progress-aware client handler
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct ProgressAwareClient {
|
||||||
|
tracker: Arc<Mutex<Option<ProgressTracker>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProgressAwareClient {
|
||||||
|
fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
tracker: Arc::new(Mutex::new(None)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn start_tracking(&self) {
|
||||||
|
if let Ok(mut tracker) = self.tracker.lock() {
|
||||||
|
*tracker = Some(ProgressTracker::new());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stop_tracking(&self) {
|
||||||
|
if let Ok(mut tracker_opt) = self.tracker.lock() {
|
||||||
|
if let Some(tracker) = tracker_opt.take() {
|
||||||
|
tracker.print_summary();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ClientHandler for ProgressAwareClient {
|
||||||
|
async fn on_progress(
|
||||||
|
&self,
|
||||||
|
params: ProgressNotificationParam,
|
||||||
|
_context: NotificationContext<RoleClient>,
|
||||||
|
) {
|
||||||
|
if let Ok(tracker_opt) = self.tracker.lock() {
|
||||||
|
if let Some(tracker) = tracker_opt.as_ref() {
|
||||||
|
tracker.handle_progress(¶ms);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn get_info(&self) -> ClientInfo {
|
||||||
|
ClientInfo {
|
||||||
|
protocol_version: Default::default(),
|
||||||
|
capabilities: ClientCapabilities::default(),
|
||||||
|
client_info: Implementation {
|
||||||
|
name: "progress-test-client".to_string(),
|
||||||
|
version: "1.0.0".to_string(),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn test_stdio_transport(records: u32) -> Result<()> {
|
||||||
|
tracing::info!("Testing STDIO Transport");
|
||||||
|
tracing::info!("================================");
|
||||||
|
|
||||||
|
let manifest_dir = env!("CARGO_MANIFEST_DIR");
|
||||||
|
let workspace_root = std::path::Path::new(manifest_dir)
|
||||||
|
.parent()
|
||||||
|
.and_then(|p| p.parent())
|
||||||
|
.ok_or_else(|| anyhow::anyhow!("Cannot find workspace root"))?;
|
||||||
|
|
||||||
|
let servers_dir = workspace_root.join("examples").join("servers");
|
||||||
|
|
||||||
|
// Start server process
|
||||||
|
let mut server_cmd = Command::new("cargo");
|
||||||
|
server_cmd
|
||||||
|
.current_dir(servers_dir)
|
||||||
|
.arg("run")
|
||||||
|
.arg("--example")
|
||||||
|
.arg("servers_progress_demo")
|
||||||
|
.arg("--")
|
||||||
|
.arg("stdio");
|
||||||
|
|
||||||
|
// Create progress-aware client handler
|
||||||
|
let client_handler = ProgressAwareClient::new();
|
||||||
|
client_handler.start_tracking();
|
||||||
|
let client_handler_clone = client_handler.clone();
|
||||||
|
|
||||||
|
let service = client_handler
|
||||||
|
.serve(TokioChildProcess::new(server_cmd)?)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// Initialize
|
||||||
|
let server_info = service.peer_info();
|
||||||
|
if let Some(info) = server_info {
|
||||||
|
tracing::info!("Connected to server: {:?}", info.server_info.name);
|
||||||
|
}
|
||||||
|
|
||||||
|
// List tools
|
||||||
|
let tools = service.list_all_tools().await?;
|
||||||
|
tracing::info!(
|
||||||
|
"Available tools: {:?}",
|
||||||
|
tools.iter().map(|t| &t.name).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
|
||||||
|
// Call stream processor tool
|
||||||
|
tracing::info!("Starting to process {} records...", records);
|
||||||
|
let tool_result = service
|
||||||
|
.call_tool(CallToolRequestParam {
|
||||||
|
name: "stream_processor".into(),
|
||||||
|
arguments: None,
|
||||||
|
})
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if let Some(content) = tool_result.content.first() {
|
||||||
|
if let Some(text) = content.as_text() {
|
||||||
|
tracing::info!("Processing completed: {}", text.text);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
service.cancel().await?;
|
||||||
|
client_handler_clone.stop_tracking();
|
||||||
|
tracing::info!("STDIO transport test completed successfully!");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
// Test SSE transport, must run the server with `cargo run --example servers_progress_demo -- sse` in the servers directory
|
||||||
|
async fn test_sse_transport(sse_url: &str, records: u32) -> Result<()> {
|
||||||
|
tracing::info!("Testing SSE Transport");
|
||||||
|
tracing::info!("=========================");
|
||||||
|
tracing::info!("SSE URL: {}", sse_url);
|
||||||
|
|
||||||
|
// Wait a bit for server to be ready
|
||||||
|
sleep(Duration::from_secs(1)).await;
|
||||||
|
|
||||||
|
let transport = SseClientTransport::start(sse_url).await?;
|
||||||
|
|
||||||
|
// Create progress-aware client handler
|
||||||
|
let client_handler = ProgressAwareClient::new();
|
||||||
|
client_handler.start_tracking();
|
||||||
|
let client_handler_clone = client_handler.clone();
|
||||||
|
|
||||||
|
let client = client_handler.serve(transport).await.inspect_err(|e| {
|
||||||
|
tracing::error!("SSE client error: {:?}", e);
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Initialize
|
||||||
|
let server_info = client.peer_info();
|
||||||
|
if let Some(info) = server_info {
|
||||||
|
tracing::info!("Connected to server: {:?}", info.server_info.name);
|
||||||
|
}
|
||||||
|
|
||||||
|
// List tools
|
||||||
|
let tools = client.list_tools(Default::default()).await?;
|
||||||
|
tracing::info!(
|
||||||
|
"Available tools: {:?}",
|
||||||
|
tools.tools.iter().map(|t| &t.name).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
|
||||||
|
// Call stream processor tool
|
||||||
|
tracing::info!("Starting to process {} records...", records);
|
||||||
|
let tool_result = client
|
||||||
|
.call_tool(CallToolRequestParam {
|
||||||
|
name: "stream_processor".into(),
|
||||||
|
arguments: None,
|
||||||
|
})
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if let Some(content) = tool_result.content.first() {
|
||||||
|
if let Some(text) = content.as_text() {
|
||||||
|
tracing::info!("Processing completed: {}", text.text);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
client.cancel().await?;
|
||||||
|
client_handler_clone.stop_tracking();
|
||||||
|
tracing::info!("SSE transport test completed successfully!");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
// Test HTTP transport, must run the server with `cargo run --example servers_progress_demo -- http` in the servers directory
|
||||||
|
async fn test_http_transport(http_url: &str, records: u32) -> Result<()> {
|
||||||
|
tracing::info!("Testing HTTP Streaming Transport");
|
||||||
|
tracing::info!("=====================================");
|
||||||
|
tracing::info!("HTTP URL: {}", http_url);
|
||||||
|
|
||||||
|
// Wait a bit for server to be ready
|
||||||
|
sleep(Duration::from_secs(1)).await;
|
||||||
|
|
||||||
|
let transport = StreamableHttpClientTransport::from_uri(http_url);
|
||||||
|
|
||||||
|
// Create progress-aware client handler
|
||||||
|
let client_handler = ProgressAwareClient::new();
|
||||||
|
client_handler.start_tracking();
|
||||||
|
let client_handler_clone = client_handler.clone();
|
||||||
|
|
||||||
|
let client = client_handler.serve(transport).await.inspect_err(|e| {
|
||||||
|
tracing::error!("HTTP client error: {:?}", e);
|
||||||
|
})?;
|
||||||
|
|
||||||
|
// Initialize
|
||||||
|
let server_info = client.peer_info();
|
||||||
|
if let Some(info) = server_info {
|
||||||
|
tracing::info!("Connected to server: {:?}", info.server_info.name);
|
||||||
|
}
|
||||||
|
|
||||||
|
// List tools
|
||||||
|
let tools = client.list_tools(Default::default()).await?;
|
||||||
|
tracing::info!(
|
||||||
|
"Available tools: {:?}",
|
||||||
|
tools.tools.iter().map(|t| &t.name).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
|
||||||
|
// Call stream processor tool
|
||||||
|
tracing::info!("Starting to process {} records...", records);
|
||||||
|
let tool_result = client
|
||||||
|
.call_tool(CallToolRequestParam {
|
||||||
|
name: "stream_processor".into(),
|
||||||
|
arguments: None,
|
||||||
|
})
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if let Some(content) = tool_result.content.first() {
|
||||||
|
if let Some(text) = content.as_text() {
|
||||||
|
tracing::info!("processing completed: {}", text.text);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
client.cancel().await?;
|
||||||
|
client_handler_clone.stop_tracking();
|
||||||
|
tracing::info!("HTTP transport test completed successfully!");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_single_test(transport_type: &TransportType, args: &Args) -> Result<bool> {
|
||||||
|
match transport_type {
|
||||||
|
TransportType::Stdio => {
|
||||||
|
test_stdio_transport(args.records).await?;
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
TransportType::Sse => match test_sse_transport(&args.sse_url, args.records).await {
|
||||||
|
Ok(_) => Ok(true),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("SSE test failed: {}", e);
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
},
|
||||||
|
TransportType::Http => match test_http_transport(&args.http_url, args.records).await {
|
||||||
|
Ok(_) => Ok(true),
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("HTTP test failed: {}", e);
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
},
|
||||||
|
TransportType::All => {
|
||||||
|
// This case is handled in main
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> Result<()> {
|
||||||
|
let args = Args::parse();
|
||||||
|
|
||||||
|
// Initialize logging
|
||||||
|
tracing_subscriber::registry()
|
||||||
|
.with(
|
||||||
|
tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
|
||||||
|
)
|
||||||
|
.with(tracing_subscriber::fmt::layer())
|
||||||
|
.init();
|
||||||
|
|
||||||
|
tracing::info!("RMCP Progress Notification Test Client");
|
||||||
|
tracing::info!("==========================================");
|
||||||
|
|
||||||
|
match &args.transport {
|
||||||
|
TransportType::All => {
|
||||||
|
// Test all transport types
|
||||||
|
let mut results = std::collections::HashMap::new();
|
||||||
|
|
||||||
|
for transport_type in [
|
||||||
|
TransportType::Stdio,
|
||||||
|
TransportType::Sse,
|
||||||
|
TransportType::Http,
|
||||||
|
] {
|
||||||
|
let transport_name = format!("{:?}", transport_type).to_uppercase();
|
||||||
|
tracing::info!("\n");
|
||||||
|
|
||||||
|
match run_single_test(&transport_type, &args).await {
|
||||||
|
Ok(success) => {
|
||||||
|
results.insert(transport_name, success);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
tracing::error!("{} test failed: {}", transport_name, e);
|
||||||
|
results.insert(transport_name, false);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
sleep(Duration::from_secs(2)).await;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Print summary
|
||||||
|
tracing::info!("\n==========================================");
|
||||||
|
tracing::info!("Test Results Summary");
|
||||||
|
tracing::info!("==========================================");
|
||||||
|
|
||||||
|
for (transport, success) in &results {
|
||||||
|
let status = if *success { "PASSED" } else { "FAILED" };
|
||||||
|
tracing::info!(" {:<10}: {}", transport, status);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
transport_type => {
|
||||||
|
run_single_test(transport_type, &args).await?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
|
|
||||||
|
|
||||||
[package]
|
[package]
|
||||||
name = "mcp-server-examples"
|
name = "mcp-server-examples"
|
||||||
version = "0.1.5"
|
version = "0.1.5"
|
||||||
|
|
@ -10,6 +8,7 @@ publish = false
|
||||||
rmcp = { workspace = true, features = [
|
rmcp = { workspace = true, features = [
|
||||||
"server",
|
"server",
|
||||||
"macros",
|
"macros",
|
||||||
|
"client",
|
||||||
"transport-sse-server",
|
"transport-sse-server",
|
||||||
"transport-io",
|
"transport-io",
|
||||||
"transport-streamable-http-server",
|
"transport-streamable-http-server",
|
||||||
|
|
@ -102,3 +101,7 @@ path = "src/elicitation_stdio.rs"
|
||||||
[[example]]
|
[[example]]
|
||||||
name = "servers_completion_stdio"
|
name = "servers_completion_stdio"
|
||||||
path = "src/completion_stdio.rs"
|
path = "src/completion_stdio.rs"
|
||||||
|
|
||||||
|
[[example]]
|
||||||
|
name = "servers_progress_demo"
|
||||||
|
path = "src/progress_demo.rs"
|
||||||
|
|
|
||||||
|
|
@ -93,6 +93,14 @@ A server demonstrating the prompt framework capabilities.
|
||||||
- Uses standard I/O transport
|
- Uses standard I/O transport
|
||||||
- Good example of prompt implementation patterns
|
- Good example of prompt implementation patterns
|
||||||
|
|
||||||
|
### Progress Demo Server (`progress_demo.rs`)
|
||||||
|
|
||||||
|
A server that demonstrates progress notifications during long-running operations.
|
||||||
|
|
||||||
|
- Provides a stream_processor tool that generates progress notifications
|
||||||
|
- Demonstrates progress notifications during long-running operations
|
||||||
|
- Can be run with `cargo run --example servers_progress_demo -- {stdio|sse|http|all}`
|
||||||
|
|
||||||
## How to Run
|
## How to Run
|
||||||
|
|
||||||
Each example can be run using Cargo:
|
Each example can be run using Cargo:
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,4 @@
|
||||||
pub mod calculator;
|
pub mod calculator;
|
||||||
pub mod counter;
|
pub mod counter;
|
||||||
pub mod generic_service;
|
pub mod generic_service;
|
||||||
|
pub mod progress_demo;
|
||||||
|
|
|
||||||
132
examples/servers/src/common/progress_demo.rs
Normal file
132
examples/servers/src/common/progress_demo.rs
Normal file
|
|
@ -0,0 +1,132 @@
|
||||||
|
use std::{
|
||||||
|
io,
|
||||||
|
pin::Pin,
|
||||||
|
task::{Context, Poll},
|
||||||
|
};
|
||||||
|
|
||||||
|
use futures::Stream;
|
||||||
|
use rmcp::{
|
||||||
|
ErrorData as McpError, RoleServer, ServerHandler, handler::server::tool::ToolRouter, model::*,
|
||||||
|
service::RequestContext, tool, tool_handler, tool_router,
|
||||||
|
};
|
||||||
|
use serde_json::json;
|
||||||
|
use tokio_stream::StreamExt;
|
||||||
|
use tracing::debug;
|
||||||
|
|
||||||
|
// a Stream data source that generates data in chunks
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct StreamDataSource {
|
||||||
|
data: Vec<u8>,
|
||||||
|
chunk_size: usize,
|
||||||
|
position: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StreamDataSource {
|
||||||
|
pub fn new(data: Vec<u8>, chunk_size: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
data,
|
||||||
|
chunk_size,
|
||||||
|
position: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
pub fn from_text(text: &str) -> Self {
|
||||||
|
Self::new(text.as_bytes().to_vec(), 1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Stream for StreamDataSource {
|
||||||
|
type Item = Result<Vec<u8>, io::Error>;
|
||||||
|
|
||||||
|
fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||||
|
let this = self.get_mut();
|
||||||
|
if this.position >= this.data.len() {
|
||||||
|
return Poll::Ready(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let start = this.position;
|
||||||
|
let end = (start + this.chunk_size).min(this.data.len());
|
||||||
|
let chunk = this.data[start..end].to_vec();
|
||||||
|
this.position = end;
|
||||||
|
Poll::Ready(Some(Ok(chunk)))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ProgressDemo {
|
||||||
|
data_source: StreamDataSource,
|
||||||
|
tool_router: ToolRouter<Self>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tool_router]
|
||||||
|
impl ProgressDemo {
|
||||||
|
#[allow(dead_code)]
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self {
|
||||||
|
tool_router: Self::tool_router(),
|
||||||
|
data_source: StreamDataSource::from_text("Hello, world!"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
#[tool(description = "Process data stream with progress updates")]
|
||||||
|
async fn stream_processor(
|
||||||
|
&self,
|
||||||
|
ctx: RequestContext<RoleServer>,
|
||||||
|
) -> Result<CallToolResult, McpError> {
|
||||||
|
let mut counter = 0;
|
||||||
|
|
||||||
|
let mut data_source = self.data_source.clone();
|
||||||
|
loop {
|
||||||
|
let chunk = data_source.next().await;
|
||||||
|
if chunk.is_none() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
|
let chunk = chunk.unwrap().unwrap();
|
||||||
|
let chunk_str = String::from_utf8_lossy(&chunk);
|
||||||
|
counter += 1;
|
||||||
|
// create progress notification param
|
||||||
|
let progress_param = ProgressNotificationParam {
|
||||||
|
progress_token: ProgressToken(NumberOrString::Number(counter)),
|
||||||
|
progress: counter as f64,
|
||||||
|
total: None,
|
||||||
|
message: Some(chunk_str.to_string()),
|
||||||
|
};
|
||||||
|
|
||||||
|
match ctx.peer.notify_progress(progress_param).await {
|
||||||
|
Ok(_) => {
|
||||||
|
debug!("Processed record: {}", chunk_str);
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
return Err(McpError::internal_error(
|
||||||
|
format!("Failed to notify progress: {}", e),
|
||||||
|
Some(json!({
|
||||||
|
"record": chunk_str,
|
||||||
|
"progress": counter,
|
||||||
|
"error": e.to_string()
|
||||||
|
})),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(CallToolResult::success(vec![Content::text(format!(
|
||||||
|
"Processed {} records successfully",
|
||||||
|
counter
|
||||||
|
))]))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tool_handler]
|
||||||
|
impl ServerHandler for ProgressDemo {
|
||||||
|
fn get_info(&self) -> ServerInfo {
|
||||||
|
ServerInfo {
|
||||||
|
protocol_version: ProtocolVersion::V_2024_11_05,
|
||||||
|
capabilities: ServerCapabilities::builder().enable_tools().build(),
|
||||||
|
server_info: Implementation::from_build_env(),
|
||||||
|
instructions: Some(
|
||||||
|
"This server demonstrates progress notifications during long-running operations. \
|
||||||
|
Use the tools to see real-time progress updates for batch processing"
|
||||||
|
.to_string(),
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
166
examples/servers/src/progress_demo.rs
Normal file
166
examples/servers/src/progress_demo.rs
Normal file
|
|
@ -0,0 +1,166 @@
|
||||||
|
use std::env;
|
||||||
|
|
||||||
|
use rmcp::{
|
||||||
|
ServiceExt,
|
||||||
|
transport::{
|
||||||
|
sse_server::{SseServer, SseServerConfig},
|
||||||
|
stdio,
|
||||||
|
streamable_http_server::{StreamableHttpService, session::local::LocalSessionManager},
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
mod common;
|
||||||
|
use common::progress_demo::ProgressDemo;
|
||||||
|
|
||||||
|
const SSE_BIND_ADDRESS: &str = "127.0.0.1:8000";
|
||||||
|
const HTTP_BIND_ADDRESS: &str = "127.0.0.1:8001";
|
||||||
|
|
||||||
|
#[tokio::main]
|
||||||
|
async fn main() -> anyhow::Result<()> {
|
||||||
|
// Get transport mode from environment variable or command line argument
|
||||||
|
let transport_mode = env::args()
|
||||||
|
.nth(1)
|
||||||
|
.unwrap_or_else(|| env::var("TRANSPORT_MODE").unwrap_or_else(|_| "stdio".to_string()));
|
||||||
|
|
||||||
|
match transport_mode.as_str() {
|
||||||
|
"stdio" => run_stdio().await,
|
||||||
|
"sse" => run_sse().await,
|
||||||
|
"http" | "streamhttp" => run_streamable_http().await,
|
||||||
|
"all" => run_all_transports().await,
|
||||||
|
_ => {
|
||||||
|
eprintln!(
|
||||||
|
"Usage: {} [stdio|sse|http|all]",
|
||||||
|
env::args().next().unwrap()
|
||||||
|
);
|
||||||
|
std::process::exit(1);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_stdio() -> anyhow::Result<()> {
|
||||||
|
let server = ProgressDemo::new();
|
||||||
|
let service = server.serve(stdio()).await.inspect_err(|e| {
|
||||||
|
tracing::error!("stdio serving error: {:?}", e);
|
||||||
|
})?;
|
||||||
|
|
||||||
|
service.waiting().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_sse() -> anyhow::Result<()> {
|
||||||
|
println!("Running SSE server");
|
||||||
|
let config = SseServerConfig {
|
||||||
|
bind: SSE_BIND_ADDRESS.parse()?,
|
||||||
|
sse_path: "/sse".to_string(),
|
||||||
|
post_path: "/message".to_string(),
|
||||||
|
ct: tokio_util::sync::CancellationToken::new(),
|
||||||
|
sse_keep_alive: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let (sse_server, router) = SseServer::new(config);
|
||||||
|
|
||||||
|
// Start the HTTP server for SSE
|
||||||
|
let listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?;
|
||||||
|
let ct = sse_server.config.ct.child_token();
|
||||||
|
|
||||||
|
let server = axum::serve(listener, router).with_graceful_shutdown(async move {
|
||||||
|
ct.cancelled().await;
|
||||||
|
tracing::info!("SSE server cancelled");
|
||||||
|
});
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Err(e) = server.await {
|
||||||
|
tracing::error!(error = %e, "SSE server shutdown with error");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Start the MCP service with SSE transport
|
||||||
|
let ct = sse_server.with_service(ProgressDemo::new);
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"Progress Demo SSE server started at http://{}/sse",
|
||||||
|
SSE_BIND_ADDRESS
|
||||||
|
);
|
||||||
|
tracing::info!("Press Ctrl+C to shutdown");
|
||||||
|
|
||||||
|
tokio::signal::ctrl_c().await?;
|
||||||
|
ct.cancel();
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_streamable_http() -> anyhow::Result<()> {
|
||||||
|
println!("Running Streamable HTTP server");
|
||||||
|
let service = StreamableHttpService::new(
|
||||||
|
|| Ok(ProgressDemo::new()),
|
||||||
|
LocalSessionManager::default().into(),
|
||||||
|
Default::default(),
|
||||||
|
);
|
||||||
|
|
||||||
|
let router = axum::Router::new().nest_service("/mcp", service);
|
||||||
|
let tcp_listener = tokio::net::TcpListener::bind(HTTP_BIND_ADDRESS).await?;
|
||||||
|
|
||||||
|
tracing::info!(
|
||||||
|
"Progress Demo HTTP server started at http://{}/mcp",
|
||||||
|
HTTP_BIND_ADDRESS
|
||||||
|
);
|
||||||
|
tracing::info!("Press Ctrl+C to shutdown");
|
||||||
|
|
||||||
|
let _ = axum::serve(tcp_listener, router)
|
||||||
|
.with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() })
|
||||||
|
.await;
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn run_all_transports() -> anyhow::Result<()> {
|
||||||
|
println!("Running all transports");
|
||||||
|
// Start SSE server
|
||||||
|
let sse_config = SseServerConfig {
|
||||||
|
bind: SSE_BIND_ADDRESS.parse()?,
|
||||||
|
sse_path: "/sse".to_string(),
|
||||||
|
post_path: "/message".to_string(),
|
||||||
|
ct: tokio_util::sync::CancellationToken::new(),
|
||||||
|
sse_keep_alive: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let (sse_server, sse_router) = SseServer::new(sse_config);
|
||||||
|
let sse_listener = tokio::net::TcpListener::bind(sse_server.config.bind).await?;
|
||||||
|
let sse_ct = sse_server.config.ct.child_token();
|
||||||
|
|
||||||
|
// Start Streamable HTTP server
|
||||||
|
let http_service = StreamableHttpService::new(
|
||||||
|
|| Ok(ProgressDemo::new()),
|
||||||
|
LocalSessionManager::default().into(),
|
||||||
|
Default::default(),
|
||||||
|
);
|
||||||
|
let http_router = axum::Router::new().nest_service("/mcp", http_service);
|
||||||
|
let http_listener = tokio::net::TcpListener::bind(HTTP_BIND_ADDRESS).await?;
|
||||||
|
|
||||||
|
// Start SSE HTTP server
|
||||||
|
let sse_http_server =
|
||||||
|
axum::serve(sse_listener, sse_router).with_graceful_shutdown(async move {
|
||||||
|
sse_ct.cancelled().await;
|
||||||
|
tracing::info!("SSE server cancelled");
|
||||||
|
});
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
if let Err(e) = sse_http_server.await {
|
||||||
|
tracing::error!(error = %e, "SSE server shutdown with error");
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// Start Streamable HTTP server
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let _ = axum::serve(http_listener, http_router)
|
||||||
|
.with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() })
|
||||||
|
.await;
|
||||||
|
});
|
||||||
|
|
||||||
|
// Start MCP service with SSE
|
||||||
|
let mcp_sse_ct = sse_server.with_service(ProgressDemo::new);
|
||||||
|
|
||||||
|
tokio::signal::ctrl_c().await?;
|
||||||
|
mcp_sse_ct.cancel();
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue