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:
parent
57d1ac94a0
commit
3dee024325
1 changed files with 52 additions and 15 deletions
|
|
@ -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(())
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue