fix(streamable-http): gracefully shutdown while client connected (#494)

* fix(streamable-http): gracefully shutdown while client connected

* fix: adviced comments

* fix: windows test build
This commit is contained in:
Jonson Petard 2025-12-02 09:22:52 +08:00 • committed by GitHub
commit 3dee024325
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,5 +1,5 @@
use rmcp::transport::streamable_http_server::{
StreamableHttpService, session::local::LocalSessionManager,
StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
};
use tracing_subscriber::{
layer::SubscriberExt,
@ -20,17 +20,24 @@ async fn main() -> anyhow::Result<()> {
)
.with(tracing_subscriber::fmt::layer())
.init();
let ct = tokio_util::sync::CancellationToken::new();
let service = StreamableHttpService::new(
|| Ok(Counter::new()),
LocalSessionManager::default().into(),
Default::default(),
StreamableHttpServerConfig {
cancellation_token: ct.child_token(),
..Default::default()
},
);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind(BIND_ADDRESS).await?;
let _ = axum::serve(tcp_listener, router)
.with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() })
.with_graceful_shutdown(async move {
tokio::signal::ctrl_c().await.unwrap();
ct.cancel();
})
.await;
Ok(())
}