chore(example): add example in progress (#271)

Signed-off-by: jokemanfire <hu.dingyang@zte.com.cn>

examples/servers/src/common/mod.rs

@@ -1,3 +1,4 @@
 pub mod calculator;
 pub mod counter;
 pub mod generic_service;
+pub mod progress_demo;

examples/servers/src/common/progress_demo.rs

@@ -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(),
+            ),
+        }
+    }
+}

examples/servers/src/progress_demo.rs

@@ -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(())
+}