refactor: provide http server as tower service (#228)

* refactor: streamable http server as tower service
This commit is contained in:
4t145 2025-05-28 19:21:40 +08:00 • committed by GitHub
commit 209be7bdb8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 1828 additions and 1180 deletions

View file

@ -2,10 +2,14 @@ use rmcp::{
ServiceExt,
service::QuitReason,
transport::{
ConfigureCommandExt, SseServer, StreamableHttpClientTransport, TokioChildProcess,
streamable_http_server::axum::StreamableHttpServer,
ConfigureCommandExt, SseServer, StreamableHttpClientTransport, StreamableHttpServerConfig,
TokioChildProcess,
streamable_http_server::{
session::local::LocalSessionManager, tower::StreamableHttpService,
},
},
};
use tokio_util::sync::CancellationToken;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
mod common;
use common::calculator::Calculator;
@ -90,10 +94,26 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> {
.wait()
.await?;
let ct = StreamableHttpServer::serve(STREAMABLE_HTTP_BIND_ADDRESS.parse()?)
.await?
.with_service(Calculator::default);
let service: StreamableHttpService<Calculator, LocalSessionManager> =
StreamableHttpService::new(
Calculator::default,
Default::default(),
StreamableHttpServerConfig {
stateful_mode: true,
sse_keep_alive: None,
},
);
let router = axum::Router::new().nest_service("/mcp", service);
let tcp_listener = tokio::net::TcpListener::bind(STREAMABLE_HTTP_BIND_ADDRESS).await?;
let ct = CancellationToken::new();
let handle = tokio::spawn({
let ct = ct.clone();
async move {
let _ = axum::serve(tcp_listener, router)
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
.await;
}
});
let exit_status = tokio::process::Command::new("node")
.arg("tests/test_with_js/streamable_client.js")
.spawn()?
@ -101,6 +121,7 @@ async fn test_with_js_streamable_http_client() -> anyhow::Result<()> {
.await?;
assert!(exit_status.success());
ct.cancel();
handle.await?;
Ok(())
}

View file

@ -1,7 +1,7 @@
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
const transport = new StreamableHTTPClientTransport(new URL(`http://127.0.0.1:8001/`));
const transport = new StreamableHTTPClientTransport(new URL(`http://127.0.0.1:8001/mcp/`));
const client = new Client(
{