Move whole rmcp crate to offcial rust sdk (#44)
* get rid of async-trait * move rmcp to official-rust-sdk * replace master with main * update readme, change license * fix typo * Update README.md Co-authored-by: Jefry Dewangga <jefrydco@gmail.com> * fix typos --------- Co-authored-by: = <=> Co-authored-by: Jefry Dewangga <jefrydco@gmail.com>
This commit is contained in:
4t145
GitHub
parent
31493cfbcc
commit
0be6067326
41 changed files with 8676 additions and 4466 deletions
97
crates/rmcp/Cargo.toml
Normal file
97
crates/rmcp/Cargo.toml
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
[package]
|
||||
name = "rmcp"
|
||||
license.workspace = true
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
repository.workspace = true
|
||||
homepage.workspace = true
|
||||
readme.workspace = true
|
||||
description = "Rust SDK for Model Context Protocol"
|
||||
documentation = "https://docs.rs/rmcp"
|
||||
|
||||
[package.metadata.docs.rs]
|
||||
all-features = true
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1.0", features = ["derive", "rc"] }
|
||||
serde_json = "1.0"
|
||||
thiserror = "2"
|
||||
chrono = { version = "0.4.38", features = ["serde"] }
|
||||
tokio = { version = "1", features = ["sync", "macros", "rt", "time"] }
|
||||
futures = "0.3"
|
||||
tracing = { version = "0.1" }
|
||||
tokio-util = { version = "0.7" }
|
||||
pin-project-lite = "0.2"
|
||||
paste = { version = "1", optional = true }
|
||||
|
||||
# for auto generate schema
|
||||
schemars = { version = "0.8", optional = true }
|
||||
|
||||
# for image encoding
|
||||
base64 = { version = "0.21", optional = true }
|
||||
|
||||
# for SSE client
|
||||
reqwest = { version = "0.12", default-features = false, features = [
|
||||
"json",
|
||||
"stream",
|
||||
"rustls-tls",
|
||||
], optional = true }
|
||||
sse-stream = { version = "0.1.3", optional = true }
|
||||
url = { version = "2.4", optional = true }
|
||||
|
||||
# For tower compatibility
|
||||
tower-service = { version = "0.3", optional = true }
|
||||
|
||||
|
||||
# for ws transport
|
||||
# tokio-tungstenite ={ version = "0.26", optional = true }
|
||||
|
||||
# for sse-server transport
|
||||
axum = { version = "0.8", features = [], optional = true }
|
||||
rand = { version = "0.9", optional = true }
|
||||
tokio-stream = { version = "0.1", optional = true }
|
||||
|
||||
# macro
|
||||
rmcp-macros = { version = "0.1", workspace = true, optional = true }
|
||||
|
||||
[features]
|
||||
default = ["base64", "macros", "server"]
|
||||
client = []
|
||||
server = ["transport-async-rw", "dep:schemars"]
|
||||
macros = ["dep:rmcp-macros", "dep:paste"]
|
||||
transport-sse = ["dep:reqwest", "dep:sse-stream", "dep:url"]
|
||||
transport-async-rw = ["tokio/io-util", "tokio-util/codec"]
|
||||
transport-io = ["transport-async-rw", "tokio/io-std"]
|
||||
transport-child-process = ["transport-async-rw", "tokio/process"]
|
||||
transport-sse-server = [
|
||||
"transport-async-rw",
|
||||
"dep:axum",
|
||||
"dep:rand",
|
||||
"dep:tokio-stream",
|
||||
]
|
||||
# transport-ws = ["transport-io", "dep:tokio-tungstenite"]
|
||||
tower = ["dep:tower-service"]
|
||||
[dev-dependencies]
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
schemars = { version = "0.8" }
|
||||
|
||||
anyhow = "1.0"
|
||||
tracing-subscriber = { version = "0.3", features = [
|
||||
"env-filter",
|
||||
"std",
|
||||
"fmt",
|
||||
] }
|
||||
[[test]]
|
||||
name = "test_tool_macros"
|
||||
required-features = ["server"]
|
||||
path = "tests/test_tool_macros.rs"
|
||||
|
||||
[[test]]
|
||||
name = "test_with_python"
|
||||
required-features = ["server", "client", "transport-sse-server", "transport-sse", "transport-child-process"]
|
||||
path = "tests/test_with_python.rs"
|
||||
|
||||
[[test]]
|
||||
name = "test_with_js"
|
||||
required-features = ["server", "transport-sse-server"]
|
||||
path = "tests/test_with_js.rs"
|
||||
183
crates/rmcp/README.md
Normal file
183
crates/rmcp/README.md
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
# RMCP
|
||||
[](https://crates.io/crates/rmcp)
|
||||
|
||||
[](https://docs.rs/rmcp/latest/rmcp)
|
||||
|
||||
A better and clean rust Model Context Protocol SDK implementation with tokio async runtime.
|
||||
|
||||
## Comparing to official SDK
|
||||
|
||||
The [Official SDK](https://github.com/modelcontextprotocol/rust-sdk/pulls) has too much limit and it was originally built for [goose](https://github.com/block/goose) rather than general using purpose.
|
||||
|
||||
All the features listed on specification would be implemented in this crate. And the first and most important thing is, this crate has the correct and intact data [types](crate::model). See it yourself.
|
||||
|
||||
## Usage
|
||||
|
||||
### Import
|
||||
```toml
|
||||
rmcp = { version = "0.1", features = ["server"] }
|
||||
```
|
||||
|
||||
### Quick start
|
||||
Start a client in one line:
|
||||
```rust,ignore
|
||||
# use rmcp::{ServiceExt, transport::child_process::TokioChildProcess};
|
||||
# use tokio::process::Command;
|
||||
|
||||
let client = ().serve(
|
||||
TokioChildProcess::new(Command::new("npx").arg("-y").arg("@modelcontextprotocol/server-everything"))?
|
||||
).await?;
|
||||
```
|
||||
|
||||
|
||||
Start a client in one line:
|
||||
```rust,ignore
|
||||
# use rmcp::{ServiceExt, transport::TokioChildProcess};
|
||||
# use tokio::process::Command;
|
||||
|
||||
let client = ().serve(
|
||||
TokioChildProcess::new(Command::new("npx").arg("-y").arg("@modelcontextprotocol/server-everything"))?
|
||||
).await?;
|
||||
```
|
||||
|
||||
|
||||
#### 1. Build a transport
|
||||
The transport type must implemented [`IntoTransport`](crate::transport::IntoTransport) trait, which allow split into a sink and a stream.
|
||||
|
||||
For client, the sink item is [`ClientJsonRpcMessage`](crate::model::ClientJsonRpcMessage) and stream item is [`ServerJsonRpcMessage`](crate::model::ServerJsonRpcMessage)
|
||||
|
||||
For server, the sink item is [`ServerJsonRpcMessage`](crate::model::ServerJsonRpcMessage) and stream item is [`ClientJsonRpcMessage`](crate::model::ClientJsonRpcMessage)
|
||||
|
||||
##### These types is automatically implemented [`IntoTransport`](crate::transport::IntoTransport) trait
|
||||
1. For type that already implement both [`Sink`](futures::Sink) and [`Stream`](futures::Stream) trait, they are automatically implemented [`IntoTransport`](crate::transport::IntoTransport) trait
|
||||
2. For tuple of sink `Tx` and stream `Rx`, type `(Tx, Rx)` are automatically implemented [`IntoTransport`](crate::transport::IntoTransport) trait
|
||||
3. For type that implement both [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`] trait, they are automatically implemented [`IntoTransport`](crate::transport::IntoTransport) trait
|
||||
4. For tuple of [`tokio::io::AsyncRead`] `R `and [`tokio::io::AsyncWrite`] `W`, type `(R, W)` are automatically implemented [`IntoTransport`](crate::transport::IntoTransport) trait
|
||||
|
||||
|
||||
```rust, ignore
|
||||
use tokio::io::{stdin, stdout};
|
||||
let transport = (stdin(), stdout());
|
||||
```
|
||||
|
||||
#### 2. Build a service
|
||||
You can easily build a service by using [`ServerHandler`](crate::handler::server) or [`ClientHandler`](crate::handler::client).
|
||||
|
||||
```rust, ignore
|
||||
let service = common::counter::Counter::new();
|
||||
```
|
||||
|
||||
Or if you want to use `tower`, you can [`TowerHandler`] as a adapter.
|
||||
|
||||
You can reference the [server examples](https://github.commodelcontextprotocol/rust-sdk/tree/release/examples/servers).
|
||||
|
||||
#### 3. Serve them together
|
||||
```rust, ignore
|
||||
// this call will finish the initialization process
|
||||
let server = service.serve(transport).await?;
|
||||
```
|
||||
|
||||
#### 4. Interact with the server
|
||||
Once the server is initialized, you can send requests or notifications:
|
||||
|
||||
```rust, ignore
|
||||
// request
|
||||
let roots = server.list_roots().await?;
|
||||
|
||||
// or send notification
|
||||
server.notify_cancelled(...).await?;
|
||||
```
|
||||
|
||||
#### 5. Waiting for service shutdown
|
||||
```rust, ignore
|
||||
let quit_reason = server.waiting().await?;
|
||||
// or cancel it
|
||||
let quit_reason = server.cancel().await?;
|
||||
```
|
||||
|
||||
### Use marcos to declaring tool
|
||||
Use `toolbox` and `tool` macros to create tool quickly.
|
||||
|
||||
Check this [file](https://github.commodelcontextprotocol/rust-sdk/tree/release/examples/servers/src/common/calculator.rs).
|
||||
```rust, ignore
|
||||
use rmcp::{ServerHandler, model::ServerInfo, schemars, tool};
|
||||
|
||||
use super::counter::Counter;
|
||||
|
||||
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
|
||||
pub struct SumRequest {
|
||||
#[schemars(description = "the left hand side number")]
|
||||
pub a: i32,
|
||||
#[schemars(description = "the right hand side number")]
|
||||
pub b: i32,
|
||||
}
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Calculator;
|
||||
|
||||
// create a static toolbox to store the tool attributes
|
||||
#[tool(tool_box)]
|
||||
impl Calculator {
|
||||
// async function
|
||||
#[tool(description = "Calculate the sum of two numbers")]
|
||||
async fn sum(&self, #[tool(aggr)] SumRequest { a, b }: SumRequest) -> String {
|
||||
(a + b).to_string()
|
||||
}
|
||||
|
||||
// sync function
|
||||
#[tool(description = "Calculate the sum of two numbers")]
|
||||
fn sub(
|
||||
&self,
|
||||
#[tool(param)]
|
||||
// this macro will transfer the schemars and serde's attributes
|
||||
#[schemars(description = "the left hand side number")]
|
||||
a: i32,
|
||||
#[tool(param)]
|
||||
#[schemars(description = "the right hand side number")]
|
||||
b: i32,
|
||||
) -> String {
|
||||
(a - b).to_string()
|
||||
}
|
||||
}
|
||||
|
||||
// impl call_tool and list_tool by querying static toolbox
|
||||
#[tool(tool_box)]
|
||||
impl ServerHandler for Calculator {
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo {
|
||||
instructions: Some("A simple calculator".into()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
```
|
||||
The only thing you should do is to make the function's return type implement `IntoCallToolResult`.
|
||||
|
||||
And you can just implement `IntoContents`, and the return value will be marked as success automatically.
|
||||
|
||||
If you return a type of `Result<T, E>` where `T` and `E` both implemented `IntoContents`, it's also OK.
|
||||
|
||||
### Manage Multi Services
|
||||
For many cases you need to manage several service in a collection, you can call `into_dyn` to convert services into the same type.
|
||||
```rust, ignore
|
||||
let service = service.into_dyn();
|
||||
```
|
||||
|
||||
|
||||
### Examples
|
||||
See [examples](https://github.commodelcontextprotocol/rust-sdk/tree/release/examples/README.md)
|
||||
|
||||
### Features
|
||||
- `client`: use client side sdk
|
||||
- `server`: use server side sdk
|
||||
- `macros`: macros default
|
||||
#### Transports
|
||||
- `transport-io`: Server stdio transport
|
||||
- `transport-sse-server`: Server SSE transport
|
||||
- `transport-child-process`: Client stdio transport
|
||||
- `transport-sse`: Client sse transport
|
||||
|
||||
## Related Resources
|
||||
- [MCP Specification](https://spec.modelcontextprotocol.io/specification/2024-11-05/)
|
||||
|
||||
- [Schema](https://github.com/modelcontextprotocol/specification/blob/main/schema/2024-11-05/schema.ts)
|
||||
17
crates/rmcp/src/error.rs
Normal file
17
crates/rmcp/src/error.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use std::fmt::Display;
|
||||
|
||||
use crate::model::ErrorData;
|
||||
|
||||
pub type Error = ErrorData;
|
||||
|
||||
impl Display for ErrorData {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}: {}", self.code.0, self.message)?;
|
||||
if let Some(data) = &self.data {
|
||||
write!(f, "({})", data)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ErrorData {}
|
||||
4
crates/rmcp/src/handler.rs
Normal file
4
crates/rmcp/src/handler.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
#[cfg(feature = "client")]
|
||||
pub mod client;
|
||||
#[cfg(feature = "server")]
|
||||
pub mod server;
|
||||
170
crates/rmcp/src/handler/client.rs
Normal file
170
crates/rmcp/src/handler/client.rs
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
use crate::error::Error as McpError;
|
||||
use crate::model::*;
|
||||
use crate::service::{Peer, RequestContext, RoleClient, Service, ServiceRole};
|
||||
|
||||
impl<H: ClientHandler> Service<RoleClient> for H {
|
||||
async fn handle_request(
|
||||
&self,
|
||||
request: <RoleClient as ServiceRole>::PeerReq,
|
||||
context: RequestContext<RoleClient>,
|
||||
) -> Result<<RoleClient as ServiceRole>::Resp, McpError> {
|
||||
match request {
|
||||
ServerRequest::PingRequest(_) => self.ping(context).await.map(ClientResult::empty),
|
||||
ServerRequest::CreateMessageRequest(request) => self
|
||||
.create_message(request.params, context)
|
||||
.await
|
||||
.map(ClientResult::CreateMessageResult),
|
||||
ServerRequest::ListRootsRequest(_) => self
|
||||
.list_roots(context)
|
||||
.await
|
||||
.map(ClientResult::ListRootsResult),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_notification(
|
||||
&self,
|
||||
notification: <RoleClient as ServiceRole>::PeerNot,
|
||||
) -> Result<(), McpError> {
|
||||
match notification {
|
||||
ServerNotification::CancelledNotification(notification) => {
|
||||
self.on_cancelled(notification.params).await
|
||||
}
|
||||
ServerNotification::ProgressNotification(notification) => {
|
||||
self.on_progress(notification.params).await
|
||||
}
|
||||
ServerNotification::LoggingMessageNotification(notification) => {
|
||||
self.on_logging_message(notification.params).await
|
||||
}
|
||||
ServerNotification::ResourceUpdatedNotification(notification) => {
|
||||
self.on_resource_updated(notification.params).await
|
||||
}
|
||||
ServerNotification::ResourceListChangedNotification(_notification_no_param) => {
|
||||
self.on_resource_list_changed().await
|
||||
}
|
||||
ServerNotification::ToolListChangedNotification(_notification_no_param) => {
|
||||
self.on_tool_list_changed().await
|
||||
}
|
||||
ServerNotification::PromptListChangedNotification(_notification_no_param) => {
|
||||
self.on_prompt_list_changed().await
|
||||
}
|
||||
};
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_peer(&self) -> Option<Peer<RoleClient>> {
|
||||
self.get_peer()
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<RoleClient>) {
|
||||
self.set_peer(peer);
|
||||
}
|
||||
|
||||
fn get_info(&self) -> <RoleClient as ServiceRole>::Info {
|
||||
self.get_info()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused_variables)]
|
||||
pub trait ClientHandler: Sized + Send + Sync + 'static {
|
||||
fn ping(
|
||||
&self,
|
||||
context: RequestContext<RoleClient>,
|
||||
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(()))
|
||||
}
|
||||
|
||||
fn create_message(
|
||||
&self,
|
||||
params: CreateMessageRequestParam,
|
||||
context: RequestContext<RoleClient>,
|
||||
) -> impl Future<Output = Result<CreateMessageResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Err(
|
||||
McpError::method_not_found::<CreateMessageRequestMethod>(),
|
||||
))
|
||||
}
|
||||
fn list_roots(
|
||||
&self,
|
||||
context: RequestContext<RoleClient>,
|
||||
) -> impl Future<Output = Result<ListRootsResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(ListRootsResult::default()))
|
||||
}
|
||||
|
||||
fn on_cancelled(
|
||||
&self,
|
||||
params: CancelledNotificationParam,
|
||||
) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_progress(
|
||||
&self,
|
||||
params: ProgressNotificationParam,
|
||||
) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_logging_message(
|
||||
&self,
|
||||
params: LoggingMessageNotificationParam,
|
||||
) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_resource_updated(
|
||||
&self,
|
||||
params: ResourceUpdatedNotificationParam,
|
||||
) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_resource_list_changed(&self) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_tool_list_changed(&self) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_prompt_list_changed(&self) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
|
||||
fn get_peer(&self) -> Option<Peer<RoleClient>>;
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<RoleClient>);
|
||||
|
||||
fn get_info(&self) -> ClientInfo {
|
||||
ClientInfo::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Do nothing, just store the peer.
|
||||
impl ClientHandler for Option<Peer<RoleClient>> {
|
||||
fn get_peer(&self) -> Option<Peer<RoleClient>> {
|
||||
self.clone()
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<RoleClient>) {
|
||||
*self = Some(peer);
|
||||
}
|
||||
}
|
||||
|
||||
/// Do nothing, even store the peer.
|
||||
impl ClientHandler for () {
|
||||
fn get_peer(&self) -> Option<Peer<RoleClient>> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<RoleClient>) {
|
||||
drop(peer);
|
||||
}
|
||||
}
|
||||
|
||||
/// Do nothing, even store the peer.
|
||||
impl ClientHandler for ClientInfo {
|
||||
fn get_peer(&self) -> Option<Peer<RoleClient>> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<RoleClient>) {
|
||||
drop(peer);
|
||||
}
|
||||
|
||||
fn get_info(&self) -> ClientInfo {
|
||||
self.clone()
|
||||
}
|
||||
}
|
||||
230
crates/rmcp/src/handler/server.rs
Normal file
230
crates/rmcp/src/handler/server.rs
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
use crate::error::Error as McpError;
|
||||
use crate::model::*;
|
||||
use crate::service::{Peer, RequestContext, RoleServer, Service, ServiceRole};
|
||||
|
||||
mod resource;
|
||||
pub mod tool;
|
||||
|
||||
impl<H: ServerHandler> Service<RoleServer> for H {
|
||||
async fn handle_request(
|
||||
&self,
|
||||
request: <RoleServer as ServiceRole>::PeerReq,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> Result<<RoleServer as ServiceRole>::Resp, McpError> {
|
||||
match request {
|
||||
ClientRequest::InitializeRequest(request) => self
|
||||
.initialize(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::InitializeResult),
|
||||
ClientRequest::PingRequest(_request) => {
|
||||
self.ping(context).await.map(ServerResult::empty)
|
||||
}
|
||||
ClientRequest::CompleteRequest(request) => self
|
||||
.complete(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::CompleteResult),
|
||||
ClientRequest::SetLevelRequest(request) => self
|
||||
.set_level(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::empty),
|
||||
ClientRequest::GetPromptRequest(request) => self
|
||||
.get_prompt(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::GetPromptResult),
|
||||
ClientRequest::ListPromptsRequest(request) => self
|
||||
.list_prompts(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::ListPromptsResult),
|
||||
ClientRequest::ListResourcesRequest(request) => self
|
||||
.list_resources(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::ListResourcesResult),
|
||||
ClientRequest::ListResourceTemplatesRequest(request) => self
|
||||
.list_resource_templates(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::ListResourceTemplatesResult),
|
||||
ClientRequest::ReadResourceRequest(request) => self
|
||||
.read_resource(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::ReadResourceResult),
|
||||
ClientRequest::SubscribeRequest(request) => self
|
||||
.subscribe(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::empty),
|
||||
ClientRequest::UnsubscribeRequest(request) => self
|
||||
.unsubscribe(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::empty),
|
||||
ClientRequest::CallToolRequest(request) => self
|
||||
.call_tool(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::CallToolResult),
|
||||
ClientRequest::ListToolsRequest(request) => self
|
||||
.list_tools(request.params, context)
|
||||
.await
|
||||
.map(ServerResult::ListToolsResult),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_notification(
|
||||
&self,
|
||||
notification: <RoleServer as ServiceRole>::PeerNot,
|
||||
) -> Result<(), McpError> {
|
||||
match notification {
|
||||
ClientNotification::CancelledNotification(notification) => {
|
||||
self.on_cancelled(notification.params).await
|
||||
}
|
||||
ClientNotification::ProgressNotification(notification) => {
|
||||
self.on_progress(notification.params).await
|
||||
}
|
||||
ClientNotification::InitializedNotification(_notification) => {
|
||||
self.on_initialized().await
|
||||
}
|
||||
ClientNotification::RootsListChangedNotification(_notification) => {
|
||||
self.on_roots_list_changed().await
|
||||
}
|
||||
};
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_peer(&self) -> Option<Peer<RoleServer>> {
|
||||
self.get_peer()
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<RoleServer>) {
|
||||
self.set_peer(peer);
|
||||
}
|
||||
|
||||
fn get_info(&self) -> <RoleServer as ServiceRole>::Info {
|
||||
self.get_info()
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused_variables)]
|
||||
pub trait ServerHandler: Sized + Clone + Send + Sync + 'static {
|
||||
fn ping(
|
||||
&self,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(()))
|
||||
}
|
||||
// handle requests
|
||||
fn initialize(
|
||||
&self,
|
||||
request: InitializeRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<InitializeResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(self.get_info()))
|
||||
}
|
||||
fn complete(
|
||||
&self,
|
||||
request: CompleteRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<CompleteResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Err(McpError::method_not_found::<CompleteRequestMethod>()))
|
||||
}
|
||||
fn set_level(
|
||||
&self,
|
||||
request: SetLevelRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
|
||||
std::future::ready(Err(McpError::method_not_found::<SetLevelRequestMethod>()))
|
||||
}
|
||||
fn get_prompt(
|
||||
&self,
|
||||
request: GetPromptRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<GetPromptResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Err(McpError::method_not_found::<GetPromptRequestMethod>()))
|
||||
}
|
||||
fn list_prompts(
|
||||
&self,
|
||||
request: PaginatedRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<ListPromptsResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(ListPromptsResult::default()))
|
||||
}
|
||||
fn list_resources(
|
||||
&self,
|
||||
request: PaginatedRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<ListResourcesResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(ListResourcesResult::default()))
|
||||
}
|
||||
fn list_resource_templates(
|
||||
&self,
|
||||
request: PaginatedRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<ListResourceTemplatesResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(ListResourceTemplatesResult::default()))
|
||||
}
|
||||
fn read_resource(
|
||||
&self,
|
||||
request: ReadResourceRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<ReadResourceResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Err(
|
||||
McpError::method_not_found::<ReadResourceRequestMethod>(),
|
||||
))
|
||||
}
|
||||
fn subscribe(
|
||||
&self,
|
||||
request: SubscribeRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
|
||||
std::future::ready(Err(McpError::method_not_found::<SubscribeRequestMethod>()))
|
||||
}
|
||||
fn unsubscribe(
|
||||
&self,
|
||||
request: UnsubscribeRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
|
||||
std::future::ready(Err(McpError::method_not_found::<UnsubscribeRequestMethod>()))
|
||||
}
|
||||
fn call_tool(
|
||||
&self,
|
||||
request: CallToolRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<CallToolResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Err(McpError::method_not_found::<CallToolRequestMethod>()))
|
||||
}
|
||||
fn list_tools(
|
||||
&self,
|
||||
request: PaginatedRequestParam,
|
||||
context: RequestContext<RoleServer>,
|
||||
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
|
||||
std::future::ready(Ok(ListToolsResult::default()))
|
||||
}
|
||||
|
||||
fn on_cancelled(
|
||||
&self,
|
||||
notification: CancelledNotificationParam,
|
||||
) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_progress(
|
||||
&self,
|
||||
notification: ProgressNotificationParam,
|
||||
) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_initialized(&self) -> impl Future<Output = ()> + Send + '_ {
|
||||
tracing::info!("client initialized");
|
||||
std::future::ready(())
|
||||
}
|
||||
fn on_roots_list_changed(&self) -> impl Future<Output = ()> + Send + '_ {
|
||||
std::future::ready(())
|
||||
}
|
||||
|
||||
fn get_peer(&self) -> Option<Peer<RoleServer>> {
|
||||
None
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<RoleServer>) {
|
||||
drop(peer);
|
||||
}
|
||||
|
||||
fn get_info(&self) -> ServerInfo {
|
||||
ServerInfo::default()
|
||||
}
|
||||
}
|
||||
1
crates/rmcp/src/handler/server/resource.rs
Normal file
1
crates/rmcp/src/handler/server/resource.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
|
||||
474
crates/rmcp/src/handler/server/tool.rs
Normal file
474
crates/rmcp/src/handler/server/tool.rs
Normal file
|
|
@ -0,0 +1,474 @@
|
|||
use std::{
|
||||
any::TypeId, borrow::Cow, collections::HashMap, future::Ready, marker::PhantomData, sync::Arc,
|
||||
};
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize, de::DeserializeOwned};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::{
|
||||
RoleServer,
|
||||
model::{CallToolRequestParam, CallToolResult, ConstString, IntoContents, JsonObject},
|
||||
service::RequestContext,
|
||||
};
|
||||
/// A shortcut for generating a JSON schema for a type.
|
||||
pub fn schema_for_type<T: JsonSchema>() -> JsonObject {
|
||||
let schema = schemars::SchemaGenerator::default().into_root_schema_for::<T>();
|
||||
let object = serde_json::to_value(schema).expect("failed to serialize schema");
|
||||
match object {
|
||||
serde_json::Value::Object(object) => object,
|
||||
_ => panic!("unexpected schema value"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Call [`schema_for_type`] with a cache
|
||||
pub fn cached_schema_for_type<T: JsonSchema + std::any::Any>() -> Arc<JsonObject> {
|
||||
thread_local! {
|
||||
static CACHE_FOR_TYPE: std::sync::RwLock<HashMap<TypeId, Arc<JsonObject>>> = Default::default();
|
||||
};
|
||||
CACHE_FOR_TYPE.with(|cache| {
|
||||
if let Some(x) = cache
|
||||
.read()
|
||||
.expect("schema cache lock poisoned")
|
||||
.get(&TypeId::of::<T>())
|
||||
{
|
||||
x.clone()
|
||||
} else {
|
||||
let schema = schema_for_type::<T>();
|
||||
let schema = Arc::new(schema);
|
||||
cache
|
||||
.write()
|
||||
.expect("schema cache lock poisoned")
|
||||
.insert(TypeId::of::<T>(), schema.clone());
|
||||
schema
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Deserialize a JSON object into a type
|
||||
pub fn parse_json_object<T: DeserializeOwned>(input: JsonObject) -> Result<T, crate::Error> {
|
||||
serde_json::from_value(serde_json::Value::Object(input)).map_err(|e| {
|
||||
crate::Error::invalid_params(
|
||||
format!("failed to deserialize parameters: {error}", error = e),
|
||||
None,
|
||||
)
|
||||
})
|
||||
}
|
||||
pub struct ToolCallContext<'service, S> {
|
||||
request_context: RequestContext<RoleServer>,
|
||||
service: &'service S,
|
||||
name: Cow<'static, str>,
|
||||
arguments: Option<JsonObject>,
|
||||
}
|
||||
|
||||
impl<'service, S> ToolCallContext<'service, S> {
|
||||
pub fn new(
|
||||
service: &'service S,
|
||||
CallToolRequestParam { name, arguments }: CallToolRequestParam,
|
||||
request_context: RequestContext<RoleServer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
request_context,
|
||||
service,
|
||||
name,
|
||||
arguments,
|
||||
}
|
||||
}
|
||||
pub fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
}
|
||||
|
||||
pub trait FromToolCallContextPart<'a, S>: Sized {
|
||||
fn from_tool_call_context_part(
|
||||
context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error>;
|
||||
}
|
||||
|
||||
pub trait IntoCallToolResult {
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::Error>;
|
||||
}
|
||||
impl IntoCallToolResult for () {
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::Error> {
|
||||
Ok(CallToolResult::success(vec![]))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: IntoContents> IntoCallToolResult for T {
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::Error> {
|
||||
Ok(CallToolResult::success(self.into_contents()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: IntoContents, E: IntoContents> IntoCallToolResult for Result<T, E> {
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::Error> {
|
||||
match self {
|
||||
Ok(value) => Ok(CallToolResult::success(value.into_contents())),
|
||||
Err(error) => Ok(CallToolResult::error(error.into_contents())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pin_project_lite::pin_project! {
|
||||
#[project = IntoCallToolResultFutProj]
|
||||
pub enum IntoCallToolResultFut<F, R> {
|
||||
Pending {
|
||||
#[pin]
|
||||
fut: F,
|
||||
_marker: PhantomData<R>,
|
||||
},
|
||||
Ready {
|
||||
#[pin]
|
||||
result: Ready<Result<CallToolResult, crate::Error>>,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<F, R> Future for IntoCallToolResultFut<F, R>
|
||||
where
|
||||
F: Future<Output = R>,
|
||||
R: IntoCallToolResult,
|
||||
{
|
||||
type Output = Result<CallToolResult, crate::Error>;
|
||||
|
||||
fn poll(
|
||||
self: std::pin::Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> std::task::Poll<Self::Output> {
|
||||
match self.project() {
|
||||
IntoCallToolResultFutProj::Pending { fut, _marker } => {
|
||||
fut.poll(cx).map(IntoCallToolResult::into_call_tool_result)
|
||||
}
|
||||
IntoCallToolResultFutProj::Ready { result } => result.poll(cx),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoCallToolResult for Result<CallToolResult, crate::Error> {
|
||||
fn into_call_tool_result(self) -> Result<CallToolResult, crate::Error> {
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
pub trait CallToolHandler<'a, S, A> {
|
||||
type Fut: Future<Output = Result<CallToolResult, crate::Error>> + Send + 'a;
|
||||
fn call(self, context: ToolCallContext<'a, S>) -> Self::Fut;
|
||||
}
|
||||
|
||||
pub type DynCallToolHandler<S> = dyn Fn(ToolCallContext<'_, S>) -> BoxFuture<'_, Result<CallToolResult, crate::Error>>
|
||||
+ Send
|
||||
+ Sync;
|
||||
/// Parameter Extractor
|
||||
pub struct Parameter<K: ConstString, V>(pub K, pub V);
|
||||
|
||||
/// Parameter Extractor
|
||||
///
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct Parameters<P>(pub P);
|
||||
|
||||
impl<P: JsonSchema> JsonSchema for Parameters<P> {
|
||||
fn schema_name() -> String {
|
||||
P::schema_name()
|
||||
}
|
||||
|
||||
fn json_schema(generator: &mut schemars::r#gen::SchemaGenerator) -> schemars::schema::Schema {
|
||||
P::json_schema(generator)
|
||||
}
|
||||
}
|
||||
|
||||
/// Callee Extractor
|
||||
pub struct Callee<'a, S>(pub &'a S);
|
||||
|
||||
impl<'a, S> FromToolCallContextPart<'a, S> for CancellationToken {
|
||||
fn from_tool_call_context_part(
|
||||
context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
|
||||
Ok((context.request_context.ct.clone(), context))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, S> FromToolCallContextPart<'a, S> for Callee<'a, S> {
|
||||
fn from_tool_call_context_part(
|
||||
context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
|
||||
Ok((Callee(context.service), context))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ToolName(pub Cow<'static, str>);
|
||||
|
||||
impl<'a, S> FromToolCallContextPart<'a, S> for ToolName {
|
||||
fn from_tool_call_context_part(
|
||||
context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
|
||||
Ok((Self(context.name.clone()), context))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, S> FromToolCallContextPart<'a, S> for &'a S {
|
||||
fn from_tool_call_context_part(
|
||||
context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
|
||||
Ok((context.service, context))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, S, K, V> FromToolCallContextPart<'a, S> for Parameter<K, V>
|
||||
where
|
||||
K: ConstString,
|
||||
V: DeserializeOwned,
|
||||
{
|
||||
fn from_tool_call_context_part(
|
||||
context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
|
||||
let arguments = context
|
||||
.arguments
|
||||
.as_ref()
|
||||
.ok_or(crate::Error::invalid_params(
|
||||
format!("missing parameter {field}", field = K::VALUE),
|
||||
None,
|
||||
))?;
|
||||
let value = arguments.get(K::VALUE).ok_or(crate::Error::invalid_params(
|
||||
format!("missing parameter {field}", field = K::VALUE),
|
||||
None,
|
||||
))?;
|
||||
let value: V = serde_json::from_value(value.clone()).map_err(|e| {
|
||||
crate::Error::invalid_params(
|
||||
format!(
|
||||
"failed to deserialize parameter {field}: {error}",
|
||||
field = K::VALUE,
|
||||
error = e
|
||||
),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
Ok((Parameter(K::default(), value), context))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, S, P> FromToolCallContextPart<'a, S> for Parameters<P>
|
||||
where
|
||||
P: DeserializeOwned,
|
||||
{
|
||||
fn from_tool_call_context_part(
|
||||
mut context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
|
||||
let arguments = context.arguments.take().unwrap_or_default();
|
||||
let value: P =
|
||||
serde_json::from_value(serde_json::Value::Object(arguments)).map_err(|e| {
|
||||
crate::Error::invalid_params(
|
||||
format!("failed to deserialize parameters: {error}", error = e),
|
||||
None,
|
||||
)
|
||||
})?;
|
||||
Ok((Parameters(value), context))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a, S> FromToolCallContextPart<'a, S> for JsonObject {
|
||||
fn from_tool_call_context_part(
|
||||
mut context: ToolCallContext<'a, S>,
|
||||
) -> Result<(Self, ToolCallContext<'a, S>), crate::Error> {
|
||||
let object = context.arguments.take().unwrap_or_default();
|
||||
Ok((object, context))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'s, S> ToolCallContext<'s, S> {
|
||||
pub fn invoke<H, A>(self, h: H) -> H::Fut
|
||||
where
|
||||
H: CallToolHandler<'s, S, A>,
|
||||
{
|
||||
h.call(self)
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub struct AsyncAdapter<P, Fut, R>(PhantomData<(fn(P) -> Fut, fn(Fut) -> R)>);
|
||||
pub struct SyncAdapter<P, R>(PhantomData<fn(P) -> R>);
|
||||
|
||||
macro_rules! impl_for {
|
||||
($($T: ident)*) => {
|
||||
impl_for!([] [$($T)*]);
|
||||
};
|
||||
// finished
|
||||
([$($Tn: ident)*] []) => {
|
||||
impl_for!(@impl $($Tn)*);
|
||||
};
|
||||
([$($Tn: ident)*] [$Tn_1: ident $($Rest: ident)*]) => {
|
||||
impl_for!(@impl $($Tn)*);
|
||||
impl_for!([$($Tn)* $Tn_1] [$($Rest)*]);
|
||||
};
|
||||
(@impl $($Tn: ident)*) => {
|
||||
impl<'s, $($Tn,)* S, F, Fut, R> CallToolHandler<'s, S, AsyncAdapter<($($Tn,)*), Fut, R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: FromToolCallContextPart<'s, S> + 's,
|
||||
)*
|
||||
F: FnOnce($($Tn,)*) -> Fut + Send + 's,
|
||||
Fut: Future<Output = R> + Send + 's,
|
||||
R: IntoCallToolResult + Send + 's,
|
||||
S: Send + Sync,
|
||||
{
|
||||
type Fut = IntoCallToolResultFut<Fut, R>;
|
||||
#[allow(unused_variables, non_snake_case)]
|
||||
fn call(
|
||||
self,
|
||||
context: ToolCallContext<'s, S>,
|
||||
) -> Self::Fut {
|
||||
$(
|
||||
let result = $Tn::from_tool_call_context_part(context);
|
||||
let ($Tn, context) = match result {
|
||||
Ok((value, context)) => (value, context),
|
||||
Err(e) => return IntoCallToolResultFut::Ready {
|
||||
result: std::future::ready(Err(e)),
|
||||
},
|
||||
};
|
||||
)*
|
||||
IntoCallToolResultFut::Pending {
|
||||
fut: self($($Tn,)*),
|
||||
_marker: PhantomData
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'s, $($Tn,)* S, F, R> CallToolHandler<'s, S, SyncAdapter<($($Tn,)*), R>> for F
|
||||
where
|
||||
$(
|
||||
$Tn: FromToolCallContextPart<'s, S> + 's,
|
||||
)*
|
||||
F: FnOnce($($Tn,)*) -> R + Send + 's,
|
||||
R: IntoCallToolResult + Send + 's,
|
||||
S: Send + Sync,
|
||||
{
|
||||
type Fut = Ready<Result<CallToolResult, crate::Error>>;
|
||||
#[allow(unused_variables, non_snake_case)]
|
||||
fn call(
|
||||
self,
|
||||
context: ToolCallContext<'s, S>,
|
||||
) -> Self::Fut {
|
||||
$(
|
||||
let result = $Tn::from_tool_call_context_part(context);
|
||||
let ($Tn, context) = match result {
|
||||
Ok((value, context)) => (value, context),
|
||||
Err(e) => return std::future::ready(Err(e)),
|
||||
};
|
||||
)*
|
||||
std::future::ready(self($($Tn,)*).into_call_tool_result())
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
impl_for!(T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 T12 T13 T14 T15);
|
||||
pub struct ToolBoxItem<S> {
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub call: Box<DynCallToolHandler<S>>,
|
||||
pub attr: crate::model::Tool,
|
||||
}
|
||||
|
||||
impl<S: Send + Sync + 'static + Clone> ToolBoxItem<S> {
|
||||
pub fn new<C>(attr: crate::model::Tool, call: C) -> Self
|
||||
where
|
||||
C: Fn(ToolCallContext<'_, S>) -> BoxFuture<'_, Result<CallToolResult, crate::Error>>
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
{
|
||||
Self {
|
||||
call: Box::new(call),
|
||||
attr,
|
||||
}
|
||||
}
|
||||
pub fn name(&self) -> &str {
|
||||
&self.attr.name
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct ToolBox<S> {
|
||||
#[allow(clippy::type_complexity)]
|
||||
pub map: std::collections::HashMap<Cow<'static, str>, ToolBoxItem<S>>,
|
||||
}
|
||||
|
||||
impl<S> ToolBox<S> {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
map: std::collections::HashMap::new(),
|
||||
}
|
||||
}
|
||||
pub fn add(&mut self, item: ToolBoxItem<S>) {
|
||||
self.map.insert(item.attr.name.clone(), item);
|
||||
}
|
||||
|
||||
pub fn remove<H, A>(&mut self, name: &str) {
|
||||
self.map.remove(name);
|
||||
}
|
||||
|
||||
pub async fn call(
|
||||
&self,
|
||||
context: ToolCallContext<'_, S>,
|
||||
) -> Result<CallToolResult, crate::Error> {
|
||||
let item = self
|
||||
.map
|
||||
.get(context.name())
|
||||
.ok_or_else(|| crate::Error::invalid_params("tool not found", None))?;
|
||||
(item.call)(context).await
|
||||
}
|
||||
|
||||
pub fn list(&self) -> Vec<crate::model::Tool> {
|
||||
self.map.values().map(|item| item.attr.clone()).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "macros")]
|
||||
#[macro_export]
|
||||
macro_rules! tool_box {
|
||||
(@pin_add $callee: ident, $attr: expr, $f: expr) => {
|
||||
$callee.add(ToolBoxItem::new($attr, |context| Box::pin($f(context))));
|
||||
};
|
||||
($server: ident { $($tool: ident),* $(,)?} ) => {
|
||||
$crate::tool_box!($server { $($tool),* } tool_box);
|
||||
};
|
||||
($server: ident { $($tool: ident),* $(,)?} $tool_box: ident) => {
|
||||
fn $tool_box() -> &'static $crate::handler::server::tool::ToolBox<$server> {
|
||||
use $crate::handler::server::tool::{ToolBox, ToolBoxItem};
|
||||
static TOOL_BOX: std::sync::OnceLock<ToolBox<$server>> = std::sync::OnceLock::new();
|
||||
TOOL_BOX.get_or_init(|| {
|
||||
let mut tool_box = ToolBox::new();
|
||||
$crate::paste!{
|
||||
$(
|
||||
$crate::tool_box!(@pin_add tool_box, $server::[< $tool _tool_attr>](), $server::[<$tool _tool_call>]);
|
||||
)*
|
||||
}
|
||||
tool_box
|
||||
})
|
||||
}
|
||||
};
|
||||
(@derive) => {
|
||||
$crate::tool_box!(@derive tool_box);
|
||||
};
|
||||
|
||||
(@derive $tool_box:ident) => {
|
||||
async fn list_tools(
|
||||
&self,
|
||||
_: $crate::model::PaginatedRequestParam,
|
||||
_: $crate::service::RequestContext<$crate::service::RoleServer>,
|
||||
) -> Result<$crate::model::ListToolsResult, $crate::Error> {
|
||||
Ok($crate::model::ListToolsResult {
|
||||
next_cursor: None,
|
||||
tools: Self::tool_box().list(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
call_tool_request_param: $crate::model::CallToolRequestParam,
|
||||
context: $crate::service::RequestContext<$crate::service::RoleServer>,
|
||||
) -> Result<$crate::model::CallToolResult, $crate::Error> {
|
||||
let context = $crate::handler::server::tool::ToolCallContext::new(self, call_tool_request_param, context);
|
||||
Self::$tool_box().call(context).await
|
||||
}
|
||||
}
|
||||
}
|
||||
35
crates/rmcp/src/lib.rs
Normal file
35
crates/rmcp/src/lib.rs
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
|
||||
mod error;
|
||||
pub use error::Error;
|
||||
|
||||
/// Basic data types in MCP specification
|
||||
pub mod model;
|
||||
#[cfg(any(feature = "client", feature = "server"))]
|
||||
pub mod service;
|
||||
#[cfg(any(feature = "client", feature = "server"))]
|
||||
pub use service::{Peer, Service, ServiceError, ServiceExt};
|
||||
#[cfg(feature = "client")]
|
||||
pub use service::{RoleClient, serve_client};
|
||||
#[cfg(feature = "server")]
|
||||
pub use service::{RoleServer, serve_server};
|
||||
|
||||
#[cfg(feature = "client")]
|
||||
pub use handler::client::ClientHandler;
|
||||
#[cfg(feature = "server")]
|
||||
pub use handler::server::ServerHandler;
|
||||
|
||||
pub mod handler;
|
||||
pub mod transport;
|
||||
|
||||
#[cfg(all(feature = "macros", feature = "server"))]
|
||||
pub use rmcp_macros::tool;
|
||||
|
||||
// re-export
|
||||
#[cfg(all(feature = "macros", feature = "server"))]
|
||||
pub use paste::paste;
|
||||
#[cfg(all(feature = "macros", feature = "server"))]
|
||||
pub use schemars;
|
||||
#[cfg(feature = "macros")]
|
||||
pub use serde;
|
||||
#[cfg(feature = "macros")]
|
||||
pub use serde_json;
|
||||
1116
crates/rmcp/src/model.rs
Normal file
1116
crates/rmcp/src/model.rs
Normal file
File diff suppressed because it is too large
Load diff
220
crates/rmcp/src/model/annotated.rs
Normal file
220
crates/rmcp/src/model/annotated.rs
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
use std::ops::{Deref, DerefMut};
|
||||
|
||||
use super::{
|
||||
RawContent, RawEmbeddedResource, RawImageContent, RawResource, RawResourceTemplate,
|
||||
RawTextContent, Role,
|
||||
};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Annotations {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub audience: Option<Vec<Role>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub priority: Option<f32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub timestamp: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
impl Annotations {
|
||||
/// Creates a new Annotations instance specifically for resources
|
||||
/// optional priority, and a timestamp (defaults to now if None)
|
||||
pub fn for_resource(priority: f32, timestamp: DateTime<Utc>) -> Self {
|
||||
assert!(
|
||||
(0.0..=1.0).contains(&priority),
|
||||
"Priority {priority} must be between 0.0 and 1.0"
|
||||
);
|
||||
Annotations {
|
||||
priority: Some(priority),
|
||||
timestamp: Some(timestamp),
|
||||
audience: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct Annotated<T: AnnotateAble> {
|
||||
#[serde(flatten)]
|
||||
pub raw: T,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub annotations: Option<Annotations>,
|
||||
}
|
||||
|
||||
impl<T: AnnotateAble> Deref for Annotated<T> {
|
||||
type Target = T;
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.raw
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AnnotateAble> DerefMut for Annotated<T> {
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
&mut self.raw
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: AnnotateAble> Annotated<T> {
|
||||
pub fn new(raw: T, annotations: Option<Annotations>) -> Self {
|
||||
Self { raw, annotations }
|
||||
}
|
||||
pub fn remove_annotation(&mut self) -> Option<Annotations> {
|
||||
self.annotations.take()
|
||||
}
|
||||
pub fn audience(&self) -> Option<&Vec<Role>> {
|
||||
self.annotations.as_ref().and_then(|a| a.audience.as_ref())
|
||||
}
|
||||
pub fn priority(&self) -> Option<f32> {
|
||||
self.annotations.as_ref().and_then(|a| a.priority)
|
||||
}
|
||||
pub fn timestamp(&self) -> Option<DateTime<Utc>> {
|
||||
self.annotations.as_ref().and_then(|a| a.timestamp)
|
||||
}
|
||||
pub fn with_audience(self, audience: Vec<Role>) -> Annotated<T>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
if let Some(annotations) = self.annotations {
|
||||
Annotated {
|
||||
raw: self.raw,
|
||||
annotations: Some(Annotations {
|
||||
audience: Some(audience),
|
||||
..annotations
|
||||
}),
|
||||
}
|
||||
} else {
|
||||
Annotated {
|
||||
raw: self.raw,
|
||||
annotations: Some(Annotations {
|
||||
audience: Some(audience),
|
||||
priority: None,
|
||||
timestamp: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn with_priority(self, priority: f32) -> Annotated<T>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
if let Some(annotations) = self.annotations {
|
||||
Annotated {
|
||||
raw: self.raw,
|
||||
annotations: Some(Annotations {
|
||||
priority: Some(priority),
|
||||
..annotations
|
||||
}),
|
||||
}
|
||||
} else {
|
||||
Annotated {
|
||||
raw: self.raw,
|
||||
annotations: Some(Annotations {
|
||||
priority: Some(priority),
|
||||
timestamp: None,
|
||||
audience: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn with_timestamp(self, timestamp: DateTime<Utc>) -> Annotated<T>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
if let Some(annotations) = self.annotations {
|
||||
Annotated {
|
||||
raw: self.raw,
|
||||
annotations: Some(Annotations {
|
||||
timestamp: Some(timestamp),
|
||||
..annotations
|
||||
}),
|
||||
}
|
||||
} else {
|
||||
Annotated {
|
||||
raw: self.raw,
|
||||
annotations: Some(Annotations {
|
||||
timestamp: Some(timestamp),
|
||||
priority: None,
|
||||
audience: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn with_timestamp_now(self) -> Annotated<T>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
self.with_timestamp(Utc::now())
|
||||
}
|
||||
}
|
||||
|
||||
mod sealed {
|
||||
pub trait Sealed {}
|
||||
}
|
||||
macro_rules! annotate {
|
||||
($T: ident) => {
|
||||
impl sealed::Sealed for $T {}
|
||||
impl AnnotateAble for $T {}
|
||||
};
|
||||
}
|
||||
|
||||
annotate!(RawContent);
|
||||
annotate!(RawTextContent);
|
||||
annotate!(RawImageContent);
|
||||
annotate!(RawEmbeddedResource);
|
||||
annotate!(RawResource);
|
||||
annotate!(RawResourceTemplate);
|
||||
|
||||
pub trait AnnotateAble: sealed::Sealed {
|
||||
fn optional_annotate(self, annotations: Option<Annotations>) -> Annotated<Self>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
Annotated::new(self, annotations)
|
||||
}
|
||||
fn annotate(self, annotations: Annotations) -> Annotated<Self>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
Annotated::new(self, Some(annotations))
|
||||
}
|
||||
fn no_annotation(self) -> Annotated<Self>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
Annotated::new(self, None)
|
||||
}
|
||||
fn with_audience(self, audience: Vec<Role>) -> Annotated<Self>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
self.annotate(Annotations {
|
||||
audience: Some(audience),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
fn with_priority(self, priority: f32) -> Annotated<Self>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
self.annotate(Annotations {
|
||||
priority: Some(priority),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
fn with_timestamp(self, timestamp: DateTime<Utc>) -> Annotated<Self>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
self.annotate(Annotations {
|
||||
timestamp: Some(timestamp),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
fn with_timestamp_now(self) -> Annotated<Self>
|
||||
where
|
||||
Self: Sized,
|
||||
{
|
||||
self.with_timestamp(Utc::now())
|
||||
}
|
||||
}
|
||||
300
crates/rmcp/src/model/capabilities.rs
Normal file
300
crates/rmcp/src/model/capabilities.rs
Normal file
|
|
@ -0,0 +1,300 @@
|
|||
use super::JsonObject;
|
||||
use paste::paste;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::BTreeMap;
|
||||
use std::marker::PhantomData;
|
||||
pub type ExperimentalCapabilities = BTreeMap<String, JsonObject>;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PromptsCapability {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub list_changed: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ResourcesCapability {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub subscribe: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub list_changed: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ToolsCapability {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub list_changed: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RootsCapabilities {
|
||||
pub list_changed: Option<bool>,
|
||||
}
|
||||
|
||||
///
|
||||
/// # Builder
|
||||
/// ```rust
|
||||
/// # use rmcp::model::ClientCapabilities;
|
||||
/// let cap = ClientCapabilities::builder()
|
||||
/// .enable_experimental()
|
||||
/// .enable_roots()
|
||||
/// .enable_roots_list_changed()
|
||||
/// .build();
|
||||
/// ```
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
|
||||
pub struct ClientCapabilities {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub experimental: Option<ExperimentalCapabilities>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub roots: Option<RootsCapabilities>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub sampling: Option<JsonObject>,
|
||||
}
|
||||
|
||||
///
|
||||
/// ## Builder
|
||||
/// ```rust
|
||||
/// # use rmcp::model::ServerCapabilities;
|
||||
/// let cap = ServerCapabilities::builder()
|
||||
/// .enable_logging()
|
||||
/// .enable_experimental()
|
||||
/// .enable_prompts()
|
||||
/// .enable_resources()
|
||||
/// .enable_tools()
|
||||
/// .enable_tool_list_changed()
|
||||
/// .build();
|
||||
/// ```
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ServerCapabilities {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub experimental: Option<ExperimentalCapabilities>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub logging: Option<JsonObject>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub prompts: Option<PromptsCapability>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub resources: Option<ResourcesCapability>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<ToolsCapability>,
|
||||
}
|
||||
|
||||
macro_rules! builder {
|
||||
($Target: ident {$($f: ident: $T: ty),* $(,)?}) => {
|
||||
paste! {
|
||||
#[derive(Default, Clone, Copy, Debug)]
|
||||
pub struct [<$Target BuilderState>]<
|
||||
$(const [<$f:upper>]: bool = false,)*
|
||||
>;
|
||||
#[derive(Debug, Default)]
|
||||
pub struct [<$Target Builder>]<S = [<$Target BuilderState>]> {
|
||||
$(pub $f: Option<$T>,)*
|
||||
pub state: PhantomData<S>
|
||||
}
|
||||
impl $Target {
|
||||
#[doc = "Create a new [`" $Target "`] builder."]
|
||||
pub fn builder() -> [<$Target Builder>] {
|
||||
<[<$Target Builder>]>::default()
|
||||
}
|
||||
}
|
||||
impl<S> [<$Target Builder>]<S> {
|
||||
pub fn build(self) -> $Target {
|
||||
$Target {
|
||||
$( $f: self.$f, )*
|
||||
}
|
||||
}
|
||||
}
|
||||
impl<S> From<[<$Target Builder>]<S>> for $Target {
|
||||
fn from(builder: [<$Target Builder>]<S>) -> Self {
|
||||
builder.build()
|
||||
}
|
||||
}
|
||||
}
|
||||
builder!($Target @toggle $($f: $T,) *);
|
||||
|
||||
};
|
||||
($Target: ident @toggle $f0: ident: $T0: ty, $($f: ident: $T: ty,)*) => {
|
||||
builder!($Target @toggle [][$f0: $T0][$($f: $T,)*]);
|
||||
};
|
||||
($Target: ident @toggle [$($ff: ident: $Tf: ty,)*][$fn: ident: $TN: ty][$fn_1: ident: $Tn_1: ty, $($ft: ident: $Tt: ty,)*]) => {
|
||||
builder!($Target @impl_toggle [$($ff: $Tf,)*][$fn: $TN][$fn_1: $Tn_1, $($ft:$Tt,)*]);
|
||||
builder!($Target @toggle [$($ff: $Tf,)* $fn: $TN,][$fn_1: $Tn_1][$($ft:$Tt,)*]);
|
||||
};
|
||||
($Target: ident @toggle [$($ff: ident: $Tf: ty,)*][$fn: ident: $TN: ty][]) => {
|
||||
builder!($Target @impl_toggle [$($ff: $Tf,)*][$fn: $TN][]);
|
||||
};
|
||||
($Target: ident @impl_toggle [$($ff: ident: $Tf: ty,)*][$fn: ident: $TN: ty][$($ft: ident: $Tt: ty,)*]) => {
|
||||
paste! {
|
||||
impl<
|
||||
$(const [<$ff:upper>]: bool,)*
|
||||
$(const [<$ft:upper>]: bool,)*
|
||||
> [<$Target Builder>]<[<$Target BuilderState>]<
|
||||
$([<$ff:upper>],)*
|
||||
false,
|
||||
$([<$ft:upper>],)*
|
||||
>> {
|
||||
pub fn [<enable_ $fn>](self) -> [<$Target Builder>]<[<$Target BuilderState>]<
|
||||
$([<$ff:upper>],)*
|
||||
true,
|
||||
$([<$ft:upper>],)*
|
||||
>> {
|
||||
[<$Target Builder>] {
|
||||
$( $ff: self.$ff, )*
|
||||
$fn: Some($TN::default()),
|
||||
$( $ft: self.$ft, )*
|
||||
state: PhantomData
|
||||
}
|
||||
}
|
||||
pub fn [<enable_ $fn _with>](self, $fn: $TN) -> [<$Target Builder>]<[<$Target BuilderState>]<
|
||||
$([<$ff:upper>],)*
|
||||
true,
|
||||
$([<$ft:upper>],)*
|
||||
>> {
|
||||
[<$Target Builder>] {
|
||||
$( $ff: self.$ff, )*
|
||||
$fn: Some($fn),
|
||||
$( $ft: self.$ft, )*
|
||||
state: PhantomData
|
||||
}
|
||||
}
|
||||
}
|
||||
// do we really need to disable some thing in builder?
|
||||
// impl<
|
||||
// $(const [<$ff:upper>]: bool,)*
|
||||
// $(const [<$ft:upper>]: bool,)*
|
||||
// > [<$Target Builder>]<[<$Target BuilderState>]<
|
||||
// $([<$ff:upper>],)*
|
||||
// true,
|
||||
// $([<$ft:upper>],)*
|
||||
// >> {
|
||||
// pub fn [<disable_ $fn>](self) -> [<$Target Builder>]<[<$Target BuilderState>]<
|
||||
// $([<$ff:upper>],)*
|
||||
// false,
|
||||
// $([<$ft:upper>],)*
|
||||
// >> {
|
||||
// [<$Target Builder>] {
|
||||
// $( $ff: self.$ff, )*
|
||||
// $fn: None,
|
||||
// $( $ft: self.$ft, )*
|
||||
// state: PhantomData
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
builder! {
|
||||
ServerCapabilities {
|
||||
experimental: ExperimentalCapabilities,
|
||||
logging: JsonObject,
|
||||
prompts: PromptsCapability,
|
||||
resources: ResourcesCapability,
|
||||
tools: ToolsCapability
|
||||
}
|
||||
}
|
||||
|
||||
impl<const E: bool, const L: bool, const P: bool, const R: bool>
|
||||
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, P, R, true>>
|
||||
{
|
||||
pub fn enable_tool_list_changed(mut self) -> Self {
|
||||
if let Some(c) = self.tools.as_mut() {
|
||||
c.list_changed = Some(true);
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<const E: bool, const L: bool, const R: bool, const T: bool>
|
||||
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, true, R, T>>
|
||||
{
|
||||
pub fn enable_prompts_list_changed(mut self) -> Self {
|
||||
if let Some(c) = self.prompts.as_mut() {
|
||||
c.list_changed = Some(true);
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl<const E: bool, const L: bool, const P: bool, const T: bool>
|
||||
ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, L, P, true, T>>
|
||||
{
|
||||
pub fn enable_resources_list_changed(mut self) -> Self {
|
||||
if let Some(c) = self.resources.as_mut() {
|
||||
c.list_changed = Some(true);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub fn enable_resources_subscribe(mut self) -> Self {
|
||||
if let Some(c) = self.resources.as_mut() {
|
||||
c.subscribe = Some(true);
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
builder! {
|
||||
ClientCapabilities{
|
||||
experimental: ExperimentalCapabilities,
|
||||
roots: RootsCapabilities,
|
||||
sampling: JsonObject,
|
||||
}
|
||||
}
|
||||
|
||||
impl<const E: bool, const S: bool>
|
||||
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, true, S>>
|
||||
{
|
||||
pub fn enable_roots_list_changed(mut self) -> Self {
|
||||
if let Some(c) = self.roots.as_mut() {
|
||||
c.list_changed = Some(true);
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn test_builder() {
|
||||
let builder = <ServerCapabilitiesBuilder>::default()
|
||||
.enable_logging()
|
||||
.enable_experimental()
|
||||
.enable_prompts()
|
||||
.enable_resources()
|
||||
.enable_tools()
|
||||
.enable_tool_list_changed();
|
||||
assert_eq!(builder.logging, Some(JsonObject::default()));
|
||||
assert_eq!(builder.prompts, Some(PromptsCapability::default()));
|
||||
assert_eq!(builder.resources, Some(ResourcesCapability::default()));
|
||||
assert_eq!(
|
||||
builder.tools,
|
||||
Some(ToolsCapability {
|
||||
list_changed: Some(true),
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
builder.experimental,
|
||||
Some(ExperimentalCapabilities::default())
|
||||
);
|
||||
let client_builder = <ClientCapabilitiesBuilder>::default()
|
||||
.enable_experimental()
|
||||
.enable_roots()
|
||||
.enable_roots_list_changed()
|
||||
.enable_sampling();
|
||||
assert_eq!(
|
||||
client_builder.experimental,
|
||||
Some(ExperimentalCapabilities::default())
|
||||
);
|
||||
assert_eq!(
|
||||
client_builder.roots,
|
||||
Some(RootsCapabilities {
|
||||
list_changed: Some(true),
|
||||
})
|
||||
);
|
||||
}
|
||||
}
|
||||
152
crates/rmcp/src/model/content.rs
Normal file
152
crates/rmcp/src/model/content.rs
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
//! Content sent around agents, extensions, and LLMs
|
||||
//! The various content types can be display to humans but also understood by models
|
||||
//! They include optional annotations used to help inform agent usage
|
||||
use super::resource::ResourceContents;
|
||||
use super::{AnnotateAble, Annotated};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RawTextContent {
|
||||
pub text: String,
|
||||
}
|
||||
pub type TextContent = Annotated<RawTextContent>;
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RawImageContent {
|
||||
/// The base64-encoded image
|
||||
pub data: String,
|
||||
pub mime_type: String,
|
||||
}
|
||||
|
||||
pub type ImageContent = Annotated<RawImageContent>;
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RawEmbeddedResource {
|
||||
pub resource: ResourceContents,
|
||||
}
|
||||
pub type EmbeddedResource = Annotated<RawEmbeddedResource>;
|
||||
|
||||
impl EmbeddedResource {
|
||||
pub fn get_text(&self) -> String {
|
||||
match &self.resource {
|
||||
ResourceContents::TextResourceContents { text, .. } => text.clone(),
|
||||
_ => String::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "camelCase")]
|
||||
pub enum RawContent {
|
||||
Text(RawTextContent),
|
||||
Image(RawImageContent),
|
||||
Resource(RawEmbeddedResource),
|
||||
}
|
||||
|
||||
pub type Content = Annotated<RawContent>;
|
||||
|
||||
impl RawContent {
|
||||
pub fn json<S: Serialize>(json: S) -> Result<Self, crate::Error> {
|
||||
let json = serde_json::to_string(&json).map_err(|e| {
|
||||
crate::Error::internal_error(
|
||||
"fail to serialize response to json",
|
||||
Some(json!(
|
||||
{"reason": e.to_string()}
|
||||
)),
|
||||
)
|
||||
})?;
|
||||
Ok(RawContent::text(json))
|
||||
}
|
||||
|
||||
pub fn text<S: Into<String>>(text: S) -> Self {
|
||||
RawContent::Text(RawTextContent { text: text.into() })
|
||||
}
|
||||
|
||||
pub fn image<S: Into<String>, T: Into<String>>(data: S, mime_type: T) -> Self {
|
||||
RawContent::Image(RawImageContent {
|
||||
data: data.into(),
|
||||
mime_type: mime_type.into(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resource(resource: ResourceContents) -> Self {
|
||||
RawContent::Resource(RawEmbeddedResource { resource })
|
||||
}
|
||||
|
||||
pub fn embedded_text<S: Into<String>, T: Into<String>>(uri: S, content: T) -> Self {
|
||||
RawContent::Resource(RawEmbeddedResource {
|
||||
resource: ResourceContents::TextResourceContents {
|
||||
uri: uri.into(),
|
||||
mime_type: Some("text".to_string()),
|
||||
text: content.into(),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the text content if this is a TextContent variant
|
||||
pub fn as_text(&self) -> Option<&RawTextContent> {
|
||||
match self {
|
||||
RawContent::Text(text) => Some(text),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the image content if this is an ImageContent variant
|
||||
pub fn as_image(&self) -> Option<&RawImageContent> {
|
||||
match self {
|
||||
RawContent::Image(image) => Some(image),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the resource content if this is an ImageContent variant
|
||||
pub fn as_resource(&self) -> Option<&RawEmbeddedResource> {
|
||||
match self {
|
||||
RawContent::Resource(resource) => Some(resource),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Content {
|
||||
pub fn text<S: Into<String>>(text: S) -> Self {
|
||||
RawContent::text(text).no_annotation()
|
||||
}
|
||||
|
||||
pub fn image<S: Into<String>, T: Into<String>>(data: S, mime_type: T) -> Self {
|
||||
RawContent::image(data, mime_type).no_annotation()
|
||||
}
|
||||
|
||||
pub fn resource(resource: ResourceContents) -> Self {
|
||||
RawContent::resource(resource).no_annotation()
|
||||
}
|
||||
|
||||
pub fn embedded_text<S: Into<String>, T: Into<String>>(uri: S, content: T) -> Self {
|
||||
RawContent::embedded_text(uri, content).no_annotation()
|
||||
}
|
||||
|
||||
pub fn json<S: Serialize>(json: S) -> Result<Self, crate::Error> {
|
||||
RawContent::json(json).map(|c| c.no_annotation())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct JsonContent<S: Serialize>(S);
|
||||
/// Types that can be converted into a list of contents
|
||||
pub trait IntoContents {
|
||||
fn into_contents(self) -> Vec<Content>;
|
||||
}
|
||||
|
||||
impl IntoContents for Content {
|
||||
fn into_contents(self) -> Vec<Content> {
|
||||
vec![self]
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoContents for String {
|
||||
fn into_contents(self) -> Vec<Content> {
|
||||
vec![Content::text(self)]
|
||||
}
|
||||
}
|
||||
162
crates/rmcp/src/model/prompt.rs
Normal file
162
crates/rmcp/src/model/prompt.rs
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
use super::content::{EmbeddedResource, ImageContent};
|
||||
use super::resource::ResourceContents;
|
||||
use super::{AnnotateAble, Annotations, RawEmbeddedResource, RawImageContent};
|
||||
use base64::engine::{Engine, general_purpose::STANDARD as BASE64_STANDARD};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// A prompt that can be used to generate text from a model
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Prompt {
|
||||
/// The name of the prompt
|
||||
pub name: String,
|
||||
/// Optional description of what the prompt does
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// Optional arguments that can be passed to customize the prompt
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub arguments: Option<Vec<PromptArgument>>,
|
||||
}
|
||||
|
||||
impl Prompt {
|
||||
/// Create a new prompt with the given name, description and arguments
|
||||
pub fn new<N, D>(
|
||||
name: N,
|
||||
description: Option<D>,
|
||||
arguments: Option<Vec<PromptArgument>>,
|
||||
) -> Self
|
||||
where
|
||||
N: Into<String>,
|
||||
D: Into<String>,
|
||||
{
|
||||
Prompt {
|
||||
name: name.into(),
|
||||
description: description.map(Into::into),
|
||||
arguments,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Represents a prompt argument that can be passed to customize the prompt
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct PromptArgument {
|
||||
/// The name of the argument
|
||||
pub name: String,
|
||||
/// A description of what the argument is used for
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// Whether this argument is required
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub required: Option<bool>,
|
||||
}
|
||||
|
||||
/// Represents the role of a message sender in a prompt conversation
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum PromptMessageRole {
|
||||
User,
|
||||
Assistant,
|
||||
}
|
||||
|
||||
/// Content types that can be included in prompt messages
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "camelCase")]
|
||||
pub enum PromptMessageContent {
|
||||
/// Plain text content
|
||||
Text { text: String },
|
||||
/// Image content with base64-encoded data
|
||||
Image {
|
||||
#[serde(flatten)]
|
||||
image: ImageContent,
|
||||
},
|
||||
/// Embedded server-side resource
|
||||
Resource { resource: EmbeddedResource },
|
||||
}
|
||||
|
||||
impl PromptMessageContent {
|
||||
pub fn text(text: impl Into<String>) -> Self {
|
||||
Self::Text { text: text.into() }
|
||||
}
|
||||
}
|
||||
|
||||
/// A message in a prompt conversation
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct PromptMessage {
|
||||
/// The role of the message sender
|
||||
pub role: PromptMessageRole,
|
||||
/// The content of the message
|
||||
pub content: PromptMessageContent,
|
||||
}
|
||||
|
||||
impl PromptMessage {
|
||||
/// Create a new text message with the given role and text content
|
||||
pub fn new_text<S: Into<String>>(role: PromptMessageRole, text: S) -> Self {
|
||||
Self {
|
||||
role,
|
||||
content: PromptMessageContent::Text { text: text.into() },
|
||||
}
|
||||
}
|
||||
#[cfg(feature = "base64")]
|
||||
pub fn new_image(
|
||||
role: PromptMessageRole,
|
||||
data: &[u8],
|
||||
mime_type: &str,
|
||||
annotations: Option<Annotations>,
|
||||
) -> Self {
|
||||
let mime_type = mime_type.into();
|
||||
|
||||
let base64 = BASE64_STANDARD.encode(data);
|
||||
|
||||
Self {
|
||||
role,
|
||||
content: PromptMessageContent::Image {
|
||||
image: RawImageContent {
|
||||
data: base64,
|
||||
mime_type,
|
||||
}
|
||||
.optional_annotate(annotations),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new resource message
|
||||
pub fn new_resource(
|
||||
role: PromptMessageRole,
|
||||
uri: String,
|
||||
mime_type: String,
|
||||
text: Option<String>,
|
||||
annotations: Option<Annotations>,
|
||||
) -> Self {
|
||||
let resource_contents = ResourceContents::TextResourceContents {
|
||||
uri,
|
||||
mime_type: Some(mime_type),
|
||||
text: text.unwrap_or_default(),
|
||||
};
|
||||
|
||||
Self {
|
||||
role,
|
||||
content: PromptMessageContent::Resource {
|
||||
resource: RawEmbeddedResource {
|
||||
resource: resource_contents,
|
||||
}
|
||||
.optional_annotate(annotations),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A template for a prompt
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct PromptTemplate {
|
||||
pub id: String,
|
||||
pub template: String,
|
||||
pub arguments: Vec<PromptArgumentTemplate>,
|
||||
}
|
||||
|
||||
/// A template for a prompt argument, this should be identical to PromptArgument
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct PromptArgumentTemplate {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub required: Option<bool>,
|
||||
}
|
||||
80
crates/rmcp/src/model/resource.rs
Normal file
80
crates/rmcp/src/model/resource.rs
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::Annotated;
|
||||
|
||||
/// Represents a resource in the extension with metadata
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RawResource {
|
||||
/// URI representing the resource location (e.g., "file:///path/to/file" or "str:///content")
|
||||
pub uri: String,
|
||||
/// Name of the resource
|
||||
pub name: String,
|
||||
/// Optional description of the resource
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
/// MIME type of the resource content ("text" or "blob")
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mime_type: Option<String>,
|
||||
|
||||
/// The size of the raw resource content, in bytes (i.e., before base64 encoding or any tokenization), if known.
|
||||
///
|
||||
/// This can be used by Hosts to display file sizes and estimate context window us
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub size: Option<u32>,
|
||||
}
|
||||
|
||||
pub type Resource = Annotated<RawResource>;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RawResourceTemplate {
|
||||
pub uri_template: String,
|
||||
pub name: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mime_type: Option<String>,
|
||||
}
|
||||
|
||||
pub type ResourceTemplate = Annotated<RawResourceTemplate>;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
|
||||
#[serde(rename_all = "camelCase", untagged)]
|
||||
pub enum ResourceContents {
|
||||
TextResourceContents {
|
||||
uri: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
mime_type: Option<String>,
|
||||
text: String,
|
||||
},
|
||||
BlobResourceContents {
|
||||
uri: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
mime_type: Option<String>,
|
||||
blob: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl ResourceContents {
|
||||
pub fn text(text: impl Into<String>, uri: impl Into<String>) -> Self {
|
||||
Self::TextResourceContents {
|
||||
uri: uri.into(),
|
||||
mime_type: Some("text".into()),
|
||||
text: text.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RawResource {
|
||||
/// Creates a new Resource from a URI with explicit mime type
|
||||
pub fn new(uri: impl Into<String>, name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
uri: uri.into(),
|
||||
name: name.into(),
|
||||
description: None,
|
||||
mime_type: None,
|
||||
size: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
41
crates/rmcp/src/model/tool.rs
Normal file
41
crates/rmcp/src/model/tool.rs
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
use std::{borrow::Cow, sync::Arc};
|
||||
|
||||
/// Tools represent a routine that a server can execute
|
||||
/// Tool calls represent requests from the client to execute one
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use super::JsonObject;
|
||||
|
||||
/// A tool that can be used by a model.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Tool {
|
||||
/// The name of the tool
|
||||
pub name: Cow<'static, str>,
|
||||
/// A description of what the tool does
|
||||
pub description: Cow<'static, str>,
|
||||
/// A JSON Schema object defining the expected parameters for the tool
|
||||
pub input_schema: Arc<JsonObject>,
|
||||
}
|
||||
|
||||
impl Tool {
|
||||
/// Create a new tool with the given name and description
|
||||
pub fn new<N, D, S>(name: N, description: D, input_schema: S) -> Self
|
||||
where
|
||||
N: Into<Cow<'static, str>>,
|
||||
D: Into<Cow<'static, str>>,
|
||||
S: Into<Arc<JsonObject>>,
|
||||
{
|
||||
Tool {
|
||||
name: name.into(),
|
||||
description: description.into(),
|
||||
input_schema: input_schema.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the schema as json value
|
||||
pub fn schema_as_json_value(&self) -> Value {
|
||||
Value::Object(self.input_schema.as_ref().clone())
|
||||
}
|
||||
}
|
||||
674
crates/rmcp/src/service.rs
Normal file
674
crates/rmcp/src/service.rs
Normal file
|
|
@ -0,0 +1,674 @@
|
|||
use crate::error::Error as McpError;
|
||||
use crate::model::{
|
||||
CancelledNotification, CancelledNotificationParam, JsonRpcMessage, Message, RequestId,
|
||||
};
|
||||
use crate::transport::IntoTransport;
|
||||
use futures::future::BoxFuture;
|
||||
use thiserror::Error;
|
||||
#[cfg(feature = "client")]
|
||||
mod client;
|
||||
#[cfg(feature = "client")]
|
||||
pub use client::*;
|
||||
#[cfg(feature = "server")]
|
||||
mod server;
|
||||
#[cfg(feature = "server")]
|
||||
pub use server::*;
|
||||
#[cfg(feature = "tower")]
|
||||
mod tower;
|
||||
#[cfg(feature = "tower")]
|
||||
pub use tower::*;
|
||||
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::instrument;
|
||||
#[derive(Error, Debug)]
|
||||
#[non_exhaustive]
|
||||
pub enum ServiceError {
|
||||
#[error("Mcp error: {0}")]
|
||||
McpError(McpError),
|
||||
#[error("Transport error: {0}")]
|
||||
Transport(std::io::Error),
|
||||
#[error("Unexpected response type")]
|
||||
UnexpectedResponse,
|
||||
#[error("task cancelled for reason {}", reason.as_deref().unwrap_or("<unknown>"))]
|
||||
Cancelled { reason: Option<String> },
|
||||
#[error("request timeout after {}", chrono::Duration::from_std(*timeout).unwrap_or_default())]
|
||||
Timeout { timeout: Duration },
|
||||
}
|
||||
|
||||
impl ServiceError {}
|
||||
trait TransferObject:
|
||||
std::fmt::Debug + Clone + serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> TransferObject for T where
|
||||
T: std::fmt::Debug
|
||||
+ serde::Serialize
|
||||
+ serde::de::DeserializeOwned
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static
|
||||
+ Clone
|
||||
{
|
||||
}
|
||||
|
||||
#[allow(private_bounds, reason = "there's no the third implementation")]
|
||||
pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
|
||||
type Req: TransferObject;
|
||||
type Resp: TransferObject;
|
||||
type Not: TryInto<CancelledNotification, Error = Self::Not>
|
||||
+ From<CancelledNotification>
|
||||
+ TransferObject;
|
||||
type PeerReq: TransferObject;
|
||||
type PeerResp: TransferObject;
|
||||
type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
|
||||
+ From<CancelledNotification>
|
||||
+ TransferObject;
|
||||
const IS_CLIENT: bool;
|
||||
type Info: TransferObject;
|
||||
type PeerInfo: TransferObject;
|
||||
}
|
||||
|
||||
pub type TxJsonRpcMessage<R> =
|
||||
JsonRpcMessage<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
|
||||
pub type RxJsonRpcMessage<R> = JsonRpcMessage<
|
||||
<R as ServiceRole>::PeerReq,
|
||||
<R as ServiceRole>::PeerResp,
|
||||
<R as ServiceRole>::PeerNot,
|
||||
>;
|
||||
|
||||
pub type TxMessage<R> =
|
||||
Message<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>;
|
||||
pub type RxMessage<R> =
|
||||
Message<<R as ServiceRole>::PeerReq, <R as ServiceRole>::PeerResp, <R as ServiceRole>::PeerNot>;
|
||||
|
||||
pub trait Service<R: ServiceRole>: Send + Sync + 'static {
|
||||
fn handle_request(
|
||||
&self,
|
||||
request: R::PeerReq,
|
||||
context: RequestContext<R>,
|
||||
) -> impl Future<Output = Result<R::Resp, McpError>> + Send + '_;
|
||||
fn handle_notification(
|
||||
&self,
|
||||
notification: R::PeerNot,
|
||||
) -> impl Future<Output = Result<(), McpError>> + Send + '_;
|
||||
fn get_peer(&self) -> Option<Peer<R>>;
|
||||
fn set_peer(&mut self, peer: Peer<R>);
|
||||
fn get_info(&self) -> R::Info;
|
||||
}
|
||||
|
||||
pub trait ServiceExt<R: ServiceRole>: Service<R> + Sized {
|
||||
/// Convert this service to a dynamic boxed service
|
||||
///
|
||||
/// This could be very helpful when you want to store the services in a collection
|
||||
fn into_dyn(self) -> Box<dyn DynService<R>> {
|
||||
Box::new(self)
|
||||
}
|
||||
fn serve<T, E, A>(
|
||||
self,
|
||||
transport: T,
|
||||
) -> impl Future<Output = Result<RunningService<R, Self>, E>> + Send
|
||||
where
|
||||
T: IntoTransport<R, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
Self: Sized,
|
||||
{
|
||||
Self::serve_with_ct(self, transport, Default::default())
|
||||
}
|
||||
fn serve_with_ct<T, E, A>(
|
||||
self,
|
||||
transport: T,
|
||||
ct: CancellationToken,
|
||||
) -> impl Future<Output = Result<RunningService<R, Self>, E>> + Send
|
||||
where
|
||||
T: IntoTransport<R, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
Self: Sized;
|
||||
}
|
||||
|
||||
impl<R: ServiceRole> Service<R> for Box<dyn DynService<R>> {
|
||||
fn handle_request(
|
||||
&self,
|
||||
request: R::PeerReq,
|
||||
context: RequestContext<R>,
|
||||
) -> impl Future<Output = Result<R::Resp, McpError>> + Send + '_ {
|
||||
DynService::handle_request(self.as_ref(), request, context)
|
||||
}
|
||||
|
||||
fn handle_notification(
|
||||
&self,
|
||||
notification: R::PeerNot,
|
||||
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
|
||||
DynService::handle_notification(self.as_ref(), notification)
|
||||
}
|
||||
|
||||
fn get_peer(&self) -> Option<Peer<R>> {
|
||||
DynService::get_peer(self.as_ref())
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<R>) {
|
||||
DynService::set_peer(self.as_mut(), peer)
|
||||
}
|
||||
|
||||
fn get_info(&self) -> R::Info {
|
||||
DynService::get_info(self.as_ref())
|
||||
}
|
||||
}
|
||||
|
||||
pub trait DynService<R: ServiceRole>: Send + Sync {
|
||||
fn handle_request(
|
||||
&self,
|
||||
request: R::PeerReq,
|
||||
context: RequestContext<R>,
|
||||
) -> BoxFuture<Result<R::Resp, McpError>>;
|
||||
fn handle_notification(&self, notification: R::PeerNot) -> BoxFuture<Result<(), McpError>>;
|
||||
fn get_peer(&self) -> Option<Peer<R>>;
|
||||
fn set_peer(&mut self, peer: Peer<R>);
|
||||
fn get_info(&self) -> R::Info;
|
||||
}
|
||||
|
||||
impl<R: ServiceRole, S: Service<R>> DynService<R> for S {
|
||||
fn handle_request(
|
||||
&self,
|
||||
request: R::PeerReq,
|
||||
context: RequestContext<R>,
|
||||
) -> BoxFuture<Result<R::Resp, McpError>> {
|
||||
Box::pin(self.handle_request(request, context))
|
||||
}
|
||||
fn handle_notification(&self, notification: R::PeerNot) -> BoxFuture<Result<(), McpError>> {
|
||||
Box::pin(self.handle_notification(notification))
|
||||
}
|
||||
fn get_peer(&self) -> Option<Peer<R>> {
|
||||
self.get_peer()
|
||||
}
|
||||
fn set_peer(&mut self, peer: Peer<R>) {
|
||||
self.set_peer(peer)
|
||||
}
|
||||
fn get_info(&self) -> R::Info {
|
||||
self.get_info()
|
||||
}
|
||||
}
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::ops::Deref;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicU32;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
pub trait RequestIdProvider: Send + Sync + 'static {
|
||||
fn next_request_id(&self) -> RequestId;
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct AtomicU32RequestIdProvider {
|
||||
id: AtomicU32,
|
||||
}
|
||||
|
||||
impl RequestIdProvider for AtomicU32RequestIdProvider {
|
||||
fn next_request_id(&self) -> RequestId {
|
||||
RequestId::Number(self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst))
|
||||
}
|
||||
}
|
||||
|
||||
type Responder<T> = tokio::sync::oneshot::Sender<T>;
|
||||
|
||||
/// A handle to a remote request
|
||||
///
|
||||
/// You can cancel it by call [`RequestHandle::cancel`] with a reason,
|
||||
///
|
||||
/// or wait for response by call [`RequestHandle::await_response`]
|
||||
#[derive(Debug)]
|
||||
pub struct RequestHandle<R: ServiceRole> {
|
||||
pub rx: tokio::sync::oneshot::Receiver<Result<R::PeerResp, ServiceError>>,
|
||||
pub options: PeerRequestOptions,
|
||||
pub peer: Peer<R>,
|
||||
pub id: RequestId,
|
||||
}
|
||||
|
||||
impl<R: ServiceRole> RequestHandle<R> {
|
||||
pub const REQUEST_TIMEOUT_REASON: &str = "request timeout";
|
||||
pub async fn await_response(self) -> Result<R::PeerResp, ServiceError> {
|
||||
if let Some(timeout) = self.options.timeout {
|
||||
let timeout_result = tokio::time::timeout(timeout, async move {
|
||||
self.rx
|
||||
.await
|
||||
.map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
|
||||
})
|
||||
.await;
|
||||
match timeout_result {
|
||||
Ok(response) => response,
|
||||
Err(_) => {
|
||||
let error = Err(ServiceError::Timeout { timeout });
|
||||
// cancel this request
|
||||
let notification = CancelledNotification {
|
||||
params: CancelledNotificationParam {
|
||||
request_id: self.id,
|
||||
reason: Some(Self::REQUEST_TIMEOUT_REASON.to_owned()),
|
||||
},
|
||||
method: crate::model::CancelledNotificationMethod,
|
||||
};
|
||||
let _ = self.peer.send_notification(notification.into()).await;
|
||||
error
|
||||
}
|
||||
}
|
||||
} else {
|
||||
self.rx
|
||||
.await
|
||||
.map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
|
||||
}
|
||||
}
|
||||
|
||||
/// Cancel this request
|
||||
pub async fn cancel(self, reason: Option<String>) -> Result<(), ServiceError> {
|
||||
let notification = CancelledNotification {
|
||||
params: CancelledNotificationParam {
|
||||
request_id: self.id,
|
||||
reason,
|
||||
},
|
||||
method: crate::model::CancelledNotificationMethod,
|
||||
};
|
||||
self.peer.send_notification(notification.into()).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum PeerSinkMessage<R: ServiceRole> {
|
||||
Request(
|
||||
R::Req,
|
||||
RequestId,
|
||||
Responder<Result<R::PeerResp, ServiceError>>,
|
||||
),
|
||||
Notification(R::Not, Responder<Result<(), ServiceError>>),
|
||||
}
|
||||
|
||||
/// An interface to fetch the remote client or server
|
||||
///
|
||||
/// For general purpose, call [`Peer::send_request`] or [`Peer::send_notification`] to send message to remote peer.
|
||||
///
|
||||
/// To create a cancellable request, call [`Peer::send_cancellable_request`].
|
||||
#[derive(Clone)]
|
||||
pub struct Peer<R: ServiceRole> {
|
||||
tx: mpsc::Sender<PeerSinkMessage<R>>,
|
||||
request_id_provider: Arc<dyn RequestIdProvider>,
|
||||
info: Arc<R::PeerInfo>,
|
||||
}
|
||||
|
||||
impl<R: ServiceRole> std::fmt::Debug for Peer<R> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("PeerSink")
|
||||
.field("tx", &self.tx)
|
||||
.field("is_client", &R::IS_CLIENT)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
type ProxyOutbound<R> = mpsc::Receiver<PeerSinkMessage<R>>;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct PeerRequestOptions {
|
||||
timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl PeerRequestOptions {
|
||||
pub fn no_options() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: ServiceRole> Peer<R> {
|
||||
const CLIENT_CHANNEL_BUFFER_SIZE: usize = 1024;
|
||||
pub fn new(
|
||||
request_id_provider: Arc<dyn RequestIdProvider>,
|
||||
peer_info: R::PeerInfo,
|
||||
) -> (Peer<R>, ProxyOutbound<R>) {
|
||||
let (tx, rx) = mpsc::channel(Self::CLIENT_CHANNEL_BUFFER_SIZE);
|
||||
(
|
||||
Self {
|
||||
tx,
|
||||
request_id_provider,
|
||||
info: peer_info.into(),
|
||||
},
|
||||
rx,
|
||||
)
|
||||
}
|
||||
pub async fn send_notification(&self, notification: R::Not) -> Result<(), ServiceError> {
|
||||
let (responder, receiver) = tokio::sync::oneshot::channel();
|
||||
self.tx
|
||||
.send(PeerSinkMessage::Notification(notification, responder))
|
||||
.await
|
||||
.map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
|
||||
receiver
|
||||
.await
|
||||
.map_err(|_e| ServiceError::Transport(std::io::Error::other("disconnected")))?
|
||||
}
|
||||
pub async fn send_request(&self, request: R::Req) -> Result<R::PeerResp, ServiceError> {
|
||||
self.send_cancellable_request(request, PeerRequestOptions::no_options())
|
||||
.await?
|
||||
.await_response()
|
||||
.await
|
||||
}
|
||||
pub async fn send_cancellable_request(
|
||||
&self,
|
||||
request: R::Req,
|
||||
options: PeerRequestOptions,
|
||||
) -> Result<RequestHandle<R>, ServiceError> {
|
||||
let id = self.request_id_provider.next_request_id();
|
||||
let (responder, receiver) = tokio::sync::oneshot::channel();
|
||||
self.tx
|
||||
.send(PeerSinkMessage::Request(request, id.clone(), responder))
|
||||
.await
|
||||
.map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
|
||||
Ok(RequestHandle {
|
||||
id,
|
||||
rx: receiver,
|
||||
options,
|
||||
peer: self.clone(),
|
||||
})
|
||||
}
|
||||
pub fn peer_info(&self) -> &R::PeerInfo {
|
||||
&self.info
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct RunningService<R: ServiceRole, S: Service<R>> {
|
||||
service: Arc<S>,
|
||||
peer: Peer<R>,
|
||||
handle: tokio::task::JoinHandle<QuitReason>,
|
||||
/// cancellation token
|
||||
ct: CancellationToken,
|
||||
}
|
||||
impl<R: ServiceRole, S: Service<R>> Deref for RunningService<R, S> {
|
||||
type Target = Peer<R>;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.peer()
|
||||
}
|
||||
}
|
||||
|
||||
impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
|
||||
#[inline]
|
||||
pub fn peer(&self) -> &Peer<R> {
|
||||
&self.peer
|
||||
}
|
||||
#[inline]
|
||||
pub fn service(&self) -> &S {
|
||||
self.service.as_ref()
|
||||
}
|
||||
pub async fn waiting(self) -> Result<QuitReason, tokio::task::JoinError> {
|
||||
self.handle.await
|
||||
}
|
||||
pub async fn cancel(self) -> Result<QuitReason, tokio::task::JoinError> {
|
||||
self.ct.cancel();
|
||||
self.waiting().await
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
|
||||
pub enum QuitReason {
|
||||
Cancelled,
|
||||
Closed,
|
||||
}
|
||||
|
||||
/// Request execution context
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RequestContext<R: ServiceRole> {
|
||||
/// this token will be cancelled when the [`CancelledNotification`] is received.
|
||||
pub ct: CancellationToken,
|
||||
pub id: RequestId,
|
||||
/// An interface to fetch the remote client or server
|
||||
pub peer: Peer<R>,
|
||||
}
|
||||
|
||||
/// Use this function to skip initialization process
|
||||
pub async fn serve_directly<R, S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
peer_info: R::PeerInfo,
|
||||
) -> Result<RunningService<R, S>, E>
|
||||
where
|
||||
R: ServiceRole,
|
||||
S: Service<R>,
|
||||
T: IntoTransport<R, E, A>,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
serve_directly_with_ct(service, transport, peer_info, Default::default()).await
|
||||
}
|
||||
|
||||
/// Use this function to skip initialization process
|
||||
pub async fn serve_directly_with_ct<R, S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
peer_info: R::PeerInfo,
|
||||
ct: CancellationToken,
|
||||
) -> Result<RunningService<R, S>, E>
|
||||
where
|
||||
R: ServiceRole,
|
||||
S: Service<R>,
|
||||
T: IntoTransport<R, E, A>,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
serve_inner(service, transport, peer_info, Default::default(), ct).await
|
||||
}
|
||||
|
||||
#[instrument(skip_all)]
|
||||
async fn serve_inner<R, S, T, E, A>(
|
||||
mut service: S,
|
||||
transport: T,
|
||||
peer_info: R::PeerInfo,
|
||||
id_provider: Arc<AtomicU32RequestIdProvider>,
|
||||
ct: CancellationToken,
|
||||
) -> Result<RunningService<R, S>, E>
|
||||
where
|
||||
R: ServiceRole,
|
||||
S: Service<R>,
|
||||
T: IntoTransport<R, E, A>,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
use futures::{SinkExt, StreamExt};
|
||||
const SINK_PROXY_BUFFER_SIZE: usize = 64;
|
||||
let (sink_proxy_tx, mut sink_proxy_rx) = tokio::sync::mpsc::channel::<
|
||||
Message<<R as ServiceRole>::Req, <R as ServiceRole>::Resp, <R as ServiceRole>::Not>,
|
||||
>(SINK_PROXY_BUFFER_SIZE);
|
||||
|
||||
if R::IS_CLIENT {
|
||||
tracing::info!(?peer_info, "Service initialized as client");
|
||||
} else {
|
||||
tracing::info!(?peer_info, "Service initialized as server");
|
||||
}
|
||||
|
||||
let (peer, mut peer_proxy) = <Peer<R>>::new(id_provider, peer_info);
|
||||
service.set_peer(peer.clone());
|
||||
let mut local_responder_pool = HashMap::new();
|
||||
let mut local_ct_pool = HashMap::<RequestId, CancellationToken>::new();
|
||||
let shared_service = Arc::new(service);
|
||||
// for return
|
||||
let service = shared_service.clone();
|
||||
|
||||
// let message_sink = tokio::sync::
|
||||
// let mut stream = std::pin::pin!(stream);
|
||||
let serve_loop_ct = ct.child_token();
|
||||
let peer_return: Peer<R> = peer.clone();
|
||||
let handle = tokio::spawn(async move {
|
||||
let (mut sink, mut stream) = transport.into_transport();
|
||||
let mut sink = std::pin::pin!(sink);
|
||||
let mut stream = std::pin::pin!(stream);
|
||||
#[derive(Debug)]
|
||||
enum Event<P, R, T> {
|
||||
ProxyMessage(P),
|
||||
PeerMessage(R),
|
||||
ToSink(T),
|
||||
}
|
||||
let quit_reason = loop {
|
||||
let evt = tokio::select! {
|
||||
m = sink_proxy_rx.recv() => {
|
||||
if let Some(m) = m {
|
||||
Event::ToSink(m)
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
}
|
||||
m = stream.next() => {
|
||||
if let Some(m) = m {
|
||||
Event::PeerMessage(m.into_message())
|
||||
} else {
|
||||
// input stream closed
|
||||
tracing::info!("input stream terminated");
|
||||
break QuitReason::Closed
|
||||
}
|
||||
}
|
||||
m = peer_proxy.recv() => {
|
||||
if let Some(m) = m {
|
||||
Event::ProxyMessage(m)
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
}
|
||||
_ = serve_loop_ct.cancelled() => {
|
||||
tracing::info!("task cancelled");
|
||||
break QuitReason::Cancelled
|
||||
}
|
||||
};
|
||||
tracing::debug!(?evt, "new event");
|
||||
match evt {
|
||||
// response and error
|
||||
Event::ToSink(e) => {
|
||||
if let Some(id) = match &e {
|
||||
Message::Response(_, id) => Some(id),
|
||||
Message::Error(_, id) => Some(id),
|
||||
_ => None,
|
||||
} {
|
||||
if let Some(ct) = local_ct_pool.remove(id) {
|
||||
ct.cancel();
|
||||
}
|
||||
let send_result = sink.send(e.into_json_rpc_message()).await;
|
||||
if let Err(error) = send_result {
|
||||
tracing::error!(%error, "fail to response message");
|
||||
}
|
||||
}
|
||||
}
|
||||
Event::ProxyMessage(PeerSinkMessage::Request(request, id, responder)) => {
|
||||
local_responder_pool.insert(id.clone(), responder);
|
||||
let send_result = sink
|
||||
.send(Message::Request(request, id.clone()).into_json_rpc_message())
|
||||
.await;
|
||||
if let Err(e) = send_result {
|
||||
if let Some(responder) = local_responder_pool.remove(&id) {
|
||||
let _ = responder
|
||||
.send(Err(ServiceError::Transport(std::io::Error::other(e))));
|
||||
}
|
||||
}
|
||||
}
|
||||
Event::ProxyMessage(PeerSinkMessage::Notification(notification, responder)) => {
|
||||
// catch cancellation notification
|
||||
let mut cancellation_param = None;
|
||||
let notification = match notification.try_into() {
|
||||
Ok::<CancelledNotification, _>(cancelled) => {
|
||||
cancellation_param.replace(cancelled.params.clone());
|
||||
cancelled.into()
|
||||
}
|
||||
Err(notification) => notification,
|
||||
};
|
||||
let send_result = sink
|
||||
.send(Message::Notification(notification).into_json_rpc_message())
|
||||
.await;
|
||||
if let Err(e) = send_result {
|
||||
let _ =
|
||||
responder.send(Err(ServiceError::Transport(std::io::Error::other(e))));
|
||||
}
|
||||
if let Some(param) = cancellation_param {
|
||||
if let Some(responder) = local_responder_pool.remove(¶m.request_id) {
|
||||
tracing::info!(id = %param.request_id, reason = param.reason, "cancelled");
|
||||
let _response_result = responder.send(Err(ServiceError::Cancelled {
|
||||
reason: param.reason.clone(),
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
Event::PeerMessage(Message::Request(request, id)) => {
|
||||
tracing::info!(%id, ?request, "received request");
|
||||
{
|
||||
let service = shared_service.clone();
|
||||
let sink = sink_proxy_tx.clone();
|
||||
let request_ct = serve_loop_ct.child_token();
|
||||
let context_ct = request_ct.child_token();
|
||||
local_ct_pool.insert(id.clone(), request_ct);
|
||||
let context = RequestContext {
|
||||
ct: context_ct,
|
||||
id: id.clone(),
|
||||
peer: peer.clone(),
|
||||
};
|
||||
tokio::spawn(async move {
|
||||
let result = service.handle_request(request, context).await;
|
||||
let response = match result {
|
||||
Ok(result) => {
|
||||
tracing::info!(%id, ?result, "response message");
|
||||
Message::Response(result, id)
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%id, ?error, "response error");
|
||||
Message::Error(error, id)
|
||||
}
|
||||
};
|
||||
let _send_result = sink.send(response).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
Event::PeerMessage(Message::Notification(notification)) => {
|
||||
tracing::info!(?notification, "received notification");
|
||||
// catch cancelled notification
|
||||
let notification = match notification.try_into() {
|
||||
Ok::<CancelledNotification, _>(cancelled) => {
|
||||
if let Some(ct) = local_ct_pool.remove(&cancelled.params.request_id) {
|
||||
tracing::info!(id = %cancelled.params.request_id, reason = cancelled.params.reason, "cancelled");
|
||||
ct.cancel();
|
||||
}
|
||||
cancelled.into()
|
||||
}
|
||||
Err(notification) => notification,
|
||||
};
|
||||
{
|
||||
let service = shared_service.clone();
|
||||
tokio::spawn(async move {
|
||||
let result = service.handle_notification(notification).await;
|
||||
if let Err(error) = result {
|
||||
tracing::warn!(%error, "Error sending notification");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
Event::PeerMessage(Message::Response(result, id)) => {
|
||||
if let Some(responder) = local_responder_pool.remove(&id) {
|
||||
let response_result = responder.send(Ok(result));
|
||||
if let Err(_error) = response_result {
|
||||
tracing::warn!(%id, "Error sending response");
|
||||
}
|
||||
}
|
||||
}
|
||||
Event::PeerMessage(Message::Error(error, id)) => {
|
||||
if let Some(responder) = local_responder_pool.remove(&id) {
|
||||
let _response_result = responder.send(Err(ServiceError::McpError(error)));
|
||||
if let Err(_error) = _response_result {
|
||||
tracing::warn!(%id, "Error sending response");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
let sink_close_result = sink.close().await;
|
||||
if let Err(e) = sink_close_result {
|
||||
tracing::error!(%e, "fail to close sink");
|
||||
}
|
||||
tracing::info!(?quit_reason, "serve finished");
|
||||
quit_reason
|
||||
});
|
||||
Ok(RunningService {
|
||||
service,
|
||||
peer: peer_return,
|
||||
handle,
|
||||
ct,
|
||||
})
|
||||
}
|
||||
287
crates/rmcp/src/service/client.rs
Normal file
287
crates/rmcp/src/service/client.rs
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
use crate::model::{
|
||||
CallToolRequest, CallToolRequestParam, CallToolResult, CancelledNotification,
|
||||
CancelledNotificationParam, ClientInfo, ClientMessage, ClientNotification, ClientRequest,
|
||||
ClientResult, CompleteRequest, CompleteRequestParam, CompleteResult, GetPromptRequest,
|
||||
GetPromptRequestParam, GetPromptResult, InitializeRequest, InitializedNotification,
|
||||
ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest,
|
||||
ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest,
|
||||
ListToolsResult, PaginatedRequestParam, PaginatedRequestParamInner, ProgressNotification,
|
||||
ProgressNotificationParam, ReadResourceRequest, ReadResourceRequestParam, ReadResourceResult,
|
||||
RootsListChangedNotification, ServerInfo, ServerNotification, ServerRequest, ServerResult,
|
||||
SetLevelRequest, SetLevelRequestParam, SubscribeRequest, SubscribeRequestParam,
|
||||
UnsubscribeRequest, UnsubscribeRequestParam,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct RoleClient;
|
||||
|
||||
impl ServiceRole for RoleClient {
|
||||
type Req = ClientRequest;
|
||||
type Resp = ClientResult;
|
||||
type Not = ClientNotification;
|
||||
type PeerReq = ServerRequest;
|
||||
type PeerResp = ServerResult;
|
||||
type PeerNot = ServerNotification;
|
||||
type Info = ClientInfo;
|
||||
type PeerInfo = ServerInfo;
|
||||
|
||||
const IS_CLIENT: bool = true;
|
||||
}
|
||||
|
||||
pub type ServerSink = Peer<RoleClient>;
|
||||
|
||||
impl<S: Service<RoleClient>> ServiceExt<RoleClient> for S {
|
||||
fn serve_with_ct<T, E, A>(
|
||||
self,
|
||||
transport: T,
|
||||
ct: CancellationToken,
|
||||
) -> impl Future<Output = Result<RunningService<RoleClient, Self>, E>> + Send
|
||||
where
|
||||
T: IntoTransport<RoleClient, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
Self: Sized,
|
||||
{
|
||||
serve_client_with_ct(self, transport, ct)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn serve_client<S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
) -> Result<RunningService<RoleClient, S>, E>
|
||||
where
|
||||
S: Service<RoleClient>,
|
||||
T: IntoTransport<RoleClient, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
{
|
||||
serve_client_with_ct(service, transport, Default::default()).await
|
||||
}
|
||||
|
||||
pub async fn serve_client_with_ct<S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
ct: CancellationToken,
|
||||
) -> Result<RunningService<RoleClient, S>, E>
|
||||
where
|
||||
S: Service<RoleClient>,
|
||||
T: IntoTransport<RoleClient, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
{
|
||||
let (sink, stream) = transport.into_transport();
|
||||
let mut sink = Box::pin(sink);
|
||||
let mut stream = Box::pin(stream);
|
||||
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
|
||||
// service
|
||||
let id = id_provider.next_request_id();
|
||||
let init_request = InitializeRequest {
|
||||
method: Default::default(),
|
||||
params: service.get_info(),
|
||||
};
|
||||
sink.send(
|
||||
ClientMessage::Request(ClientRequest::InitializeRequest(init_request), id.clone())
|
||||
.into_json_rpc_message(),
|
||||
)
|
||||
.await?;
|
||||
let (response, response_id) = stream
|
||||
.next()
|
||||
.await
|
||||
.ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
"expect initialize response",
|
||||
))?
|
||||
.into_message()
|
||||
.into_result()
|
||||
.ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"expect initialize result",
|
||||
))?;
|
||||
if id != response_id {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"conflict initialize response id",
|
||||
)
|
||||
.into());
|
||||
}
|
||||
let response = response.map_err(std::io::Error::other)?;
|
||||
let ServerResult::InitializeResult(initialize_result) = response else {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"expect initialize result",
|
||||
)
|
||||
.into());
|
||||
};
|
||||
// send notification
|
||||
let notification = ClientMessage::Notification(ClientNotification::InitializedNotification(
|
||||
InitializedNotification {
|
||||
method: Default::default(),
|
||||
},
|
||||
));
|
||||
sink.send(notification.into_json_rpc_message()).await?;
|
||||
serve_inner(service, (sink, stream), initialize_result, id_provider, ct).await
|
||||
}
|
||||
|
||||
macro_rules! method {
|
||||
(peer_req $method:ident $Req:ident() => $Resp: ident ) => {
|
||||
pub async fn $method(&self) -> Result<$Resp, ServiceError> {
|
||||
let result = self
|
||||
.send_request(ClientRequest::$Req($Req {
|
||||
method: Default::default(),
|
||||
}))
|
||||
.await?;
|
||||
match result {
|
||||
ServerResult::$Resp(result) => Ok(result),
|
||||
_ => Err(ServiceError::UnexpectedResponse),
|
||||
}
|
||||
}
|
||||
};
|
||||
(peer_req $method:ident $Req:ident($Param: ident) => $Resp: ident ) => {
|
||||
pub async fn $method(&self, params: $Param) -> Result<$Resp, ServiceError> {
|
||||
let result = self
|
||||
.send_request(ClientRequest::$Req($Req {
|
||||
method: Default::default(),
|
||||
params,
|
||||
}))
|
||||
.await?;
|
||||
match result {
|
||||
ServerResult::$Resp(result) => Ok(result),
|
||||
_ => Err(ServiceError::UnexpectedResponse),
|
||||
}
|
||||
}
|
||||
};
|
||||
(peer_req $method:ident $Req:ident($Param: ident)) => {
|
||||
pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
|
||||
let result = self
|
||||
.send_request(ClientRequest::$Req($Req {
|
||||
method: Default::default(),
|
||||
params,
|
||||
}))
|
||||
.await?;
|
||||
match result {
|
||||
ServerResult::EmptyResult(_) => Ok(()),
|
||||
_ => Err(ServiceError::UnexpectedResponse),
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
(peer_not $method:ident $Not:ident($Param: ident)) => {
|
||||
pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
|
||||
self.send_notification(ClientNotification::$Not($Not {
|
||||
method: Default::default(),
|
||||
params,
|
||||
}))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
(peer_not $method:ident $Not:ident) => {
|
||||
pub async fn $method(&self) -> Result<(), ServiceError> {
|
||||
self.send_notification(ClientNotification::$Not($Not {
|
||||
method: Default::default(),
|
||||
}))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
impl Peer<RoleClient> {
|
||||
method!(peer_req complete CompleteRequest(CompleteRequestParam) => CompleteResult);
|
||||
method!(peer_req set_level SetLevelRequest(SetLevelRequestParam));
|
||||
method!(peer_req get_prompt GetPromptRequest(GetPromptRequestParam) => GetPromptResult);
|
||||
method!(peer_req list_prompts ListPromptsRequest(PaginatedRequestParam) => ListPromptsResult);
|
||||
method!(peer_req list_resources ListResourcesRequest(PaginatedRequestParam) => ListResourcesResult);
|
||||
method!(peer_req list_resource_templates ListResourceTemplatesRequest(PaginatedRequestParam) => ListResourceTemplatesResult);
|
||||
method!(peer_req read_resource ReadResourceRequest(ReadResourceRequestParam) => ReadResourceResult);
|
||||
method!(peer_req subscribe SubscribeRequest(SubscribeRequestParam) );
|
||||
method!(peer_req unsubscribe UnsubscribeRequest(UnsubscribeRequestParam));
|
||||
method!(peer_req call_tool CallToolRequest(CallToolRequestParam) => CallToolResult);
|
||||
method!(peer_req list_tools ListToolsRequest(PaginatedRequestParam) => ListToolsResult);
|
||||
|
||||
method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam));
|
||||
method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam));
|
||||
method!(peer_not notify_initialized InitializedNotification);
|
||||
method!(peer_not notify_roots_list_changed RootsListChangedNotification);
|
||||
}
|
||||
|
||||
impl Peer<RoleClient> {
|
||||
/// A wrapper method for [`Peer<RoleClient>::list_tools`].
|
||||
///
|
||||
/// This function will call [`Peer<RoleClient>::list_tools`] multiple times until all tools are listed.
|
||||
pub async fn list_all_tools(&self) -> Result<Vec<crate::model::Tool>, ServiceError> {
|
||||
let mut tools = Vec::new();
|
||||
let mut cursor = None;
|
||||
loop {
|
||||
let result = self
|
||||
.list_tools(Some(PaginatedRequestParamInner { cursor }))
|
||||
.await?;
|
||||
tools.extend(result.tools);
|
||||
cursor = result.next_cursor;
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(tools)
|
||||
}
|
||||
|
||||
/// A wrapper method for [`Peer<RoleClient>::list_resources`].
|
||||
///
|
||||
/// This function will call [`Peer<RoleClient>::list_resources`] multiple times until all resources are listed.
|
||||
pub async fn list_all_prompts(&self) -> Result<Vec<crate::model::Prompt>, ServiceError> {
|
||||
let mut prompts = Vec::new();
|
||||
let mut cursor = None;
|
||||
loop {
|
||||
let result = self
|
||||
.list_prompts(Some(PaginatedRequestParamInner { cursor }))
|
||||
.await?;
|
||||
prompts.extend(result.prompts);
|
||||
cursor = result.next_cursor;
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(prompts)
|
||||
}
|
||||
|
||||
/// A wrapper method for [`Peer<RoleClient>::list_resources`].
|
||||
///
|
||||
/// This function will call [`Peer<RoleClient>::list_resources`] multiple times until all resources are listed.
|
||||
pub async fn list_all_resources(&self) -> Result<Vec<crate::model::Resource>, ServiceError> {
|
||||
let mut resources = Vec::new();
|
||||
let mut cursor = None;
|
||||
loop {
|
||||
let result = self
|
||||
.list_resources(Some(PaginatedRequestParamInner { cursor }))
|
||||
.await?;
|
||||
resources.extend(result.resources);
|
||||
cursor = result.next_cursor;
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(resources)
|
||||
}
|
||||
|
||||
/// A wrapper method for [`Peer<RoleClient>::list_resource_templates`].
|
||||
///
|
||||
/// This function will call [`Peer<RoleClient>::list_resource_templates`] multiple times until all resource templates are listed.
|
||||
pub async fn list_all_resource_templates(
|
||||
&self,
|
||||
) -> Result<Vec<crate::model::ResourceTemplate>, ServiceError> {
|
||||
let mut resource_templates = Vec::new();
|
||||
let mut cursor = None;
|
||||
loop {
|
||||
let result = self
|
||||
.list_resource_templates(Some(PaginatedRequestParamInner { cursor }))
|
||||
.await?;
|
||||
resource_templates.extend(result.resource_templates);
|
||||
cursor = result.next_cursor;
|
||||
if cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(resource_templates)
|
||||
}
|
||||
}
|
||||
204
crates/rmcp/src/service/server.rs
Normal file
204
crates/rmcp/src/service/server.rs
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
use crate::model::{
|
||||
CancelledNotification, CancelledNotificationParam, ClientInfo, ClientNotification,
|
||||
ClientRequest, ClientResult, CreateMessageRequest, CreateMessageRequestParam,
|
||||
CreateMessageResult, ListRootsRequest, ListRootsResult, LoggingMessageNotification,
|
||||
LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam,
|
||||
PromptListChangedNotification, ResourceListChangedNotification, ResourceUpdatedNotification,
|
||||
ResourceUpdatedNotificationParam, ServerInfo, ServerMessage, ServerNotification, ServerRequest,
|
||||
ServerResult, ToolListChangedNotification,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct RoleServer;
|
||||
|
||||
impl ServiceRole for RoleServer {
|
||||
type Req = ServerRequest;
|
||||
type Resp = ServerResult;
|
||||
type Not = ServerNotification;
|
||||
type PeerReq = ClientRequest;
|
||||
type PeerResp = ClientResult;
|
||||
type PeerNot = ClientNotification;
|
||||
type Info = ServerInfo;
|
||||
type PeerInfo = ClientInfo;
|
||||
const IS_CLIENT: bool = false;
|
||||
}
|
||||
|
||||
pub type ClientSink = Peer<RoleServer>;
|
||||
|
||||
impl<S: Service<RoleServer>> ServiceExt<RoleServer> for S {
|
||||
fn serve_with_ct<T, E, A>(
|
||||
self,
|
||||
transport: T,
|
||||
ct: CancellationToken,
|
||||
) -> impl Future<Output = Result<RunningService<RoleServer, Self>, E>> + Send
|
||||
where
|
||||
T: IntoTransport<RoleServer, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
Self: Sized,
|
||||
{
|
||||
serve_server_with_ct(self, transport, ct)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn serve_server<S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
) -> Result<RunningService<RoleServer, S>, E>
|
||||
where
|
||||
S: Service<RoleServer>,
|
||||
T: IntoTransport<RoleServer, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
{
|
||||
serve_server_with_ct(service, transport, CancellationToken::new()).await
|
||||
}
|
||||
|
||||
pub async fn serve_server_with_ct<S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
ct: CancellationToken,
|
||||
) -> Result<RunningService<RoleServer, S>, E>
|
||||
where
|
||||
S: Service<RoleServer>,
|
||||
T: IntoTransport<RoleServer, E, A>,
|
||||
E: std::error::Error + From<std::io::Error> + Send + Sync + 'static,
|
||||
{
|
||||
let (sink, stream) = transport.into_transport();
|
||||
let mut sink = Box::pin(sink);
|
||||
let mut stream = Box::pin(stream);
|
||||
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
|
||||
|
||||
// service
|
||||
let (request, id) = stream
|
||||
.next()
|
||||
.await
|
||||
.ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
"expect initialize request",
|
||||
))?
|
||||
.into_message()
|
||||
.into_request()
|
||||
.ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"expect initialize request",
|
||||
))?;
|
||||
let ClientRequest::InitializeRequest(peer_info) = request else {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"expect initialize request",
|
||||
)
|
||||
.into());
|
||||
};
|
||||
let init_response = service.get_info();
|
||||
sink.send(
|
||||
ServerMessage::Response(ServerResult::InitializeResult(init_response), id)
|
||||
.into_json_rpc_message(),
|
||||
)
|
||||
.await?;
|
||||
// waiting for notification
|
||||
let notification = stream
|
||||
.next()
|
||||
.await
|
||||
.ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
"expect initialize notification",
|
||||
))?
|
||||
.into_message()
|
||||
.into_notification()
|
||||
.ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"expect initialize notification",
|
||||
))?;
|
||||
let ClientNotification::InitializedNotification(_) = notification else {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"expect initialize notification",
|
||||
)
|
||||
.into());
|
||||
};
|
||||
serve_inner(service, (sink, stream), peer_info.params, id_provider, ct).await
|
||||
}
|
||||
|
||||
macro_rules! method {
|
||||
(peer_req $method:ident $Req:ident() => $Resp: ident ) => {
|
||||
pub async fn $method(&self) -> Result<$Resp, ServiceError> {
|
||||
let result = self
|
||||
.send_request(ServerRequest::$Req($Req {
|
||||
method: Default::default(),
|
||||
}))
|
||||
.await?;
|
||||
match result {
|
||||
ClientResult::$Resp(result) => Ok(result),
|
||||
_ => Err(ServiceError::UnexpectedResponse),
|
||||
}
|
||||
}
|
||||
};
|
||||
(peer_req $method:ident $Req:ident($Param: ident) => $Resp: ident ) => {
|
||||
pub async fn $method(&self, params: $Param) -> Result<$Resp, ServiceError> {
|
||||
let result = self
|
||||
.send_request(ServerRequest::$Req($Req {
|
||||
method: Default::default(),
|
||||
params,
|
||||
}))
|
||||
.await?;
|
||||
match result {
|
||||
ClientResult::$Resp(result) => Ok(result),
|
||||
_ => Err(ServiceError::UnexpectedResponse),
|
||||
}
|
||||
}
|
||||
};
|
||||
(peer_req $method:ident $Req:ident($Param: ident)) => {
|
||||
pub fn $method(
|
||||
&self,
|
||||
params: $Param,
|
||||
) -> impl Future<Output = Result<(), ServiceError>> + Send + '_ {
|
||||
async move {
|
||||
let result = self
|
||||
.send_request(ServerRequest::$Req($Req {
|
||||
method: Default::default(),
|
||||
params,
|
||||
}))
|
||||
.await?;
|
||||
match result {
|
||||
ClientResult::EmptyResult(_) => Ok(()),
|
||||
_ => Err(ServiceError::UnexpectedResponse),
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
(peer_not $method:ident $Not:ident($Param: ident)) => {
|
||||
pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
|
||||
self.send_notification(ServerNotification::$Not($Not {
|
||||
method: Default::default(),
|
||||
params,
|
||||
}))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
(peer_not $method:ident $Not:ident) => {
|
||||
pub async fn $method(&self) -> Result<(), ServiceError> {
|
||||
self.send_notification(ServerNotification::$Not($Not {
|
||||
method: Default::default(),
|
||||
}))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
impl Peer<RoleServer> {
|
||||
method!(peer_req create_message CreateMessageRequest(CreateMessageRequestParam) => CreateMessageResult);
|
||||
method!(peer_req list_roots ListRootsRequest() => ListRootsResult);
|
||||
|
||||
method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam));
|
||||
method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam));
|
||||
method!(peer_not notify_logging_message LoggingMessageNotification(LoggingMessageNotificationParam));
|
||||
method!(peer_not notify_resource_updated ResourceUpdatedNotification(ResourceUpdatedNotificationParam));
|
||||
method!(peer_not notify_resource_list_changed ResourceListChangedNotification);
|
||||
method!(peer_not notify_tool_list_changed ToolListChangedNotification);
|
||||
method!(peer_not notify_prompt_list_changed PromptListChangedNotification);
|
||||
}
|
||||
62
crates/rmcp/src/service/tower.rs
Normal file
62
crates/rmcp/src/service/tower.rs
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
use std::{future::poll_fn, marker::PhantomData};
|
||||
|
||||
use tower_service::Service as TowerService;
|
||||
|
||||
use crate::service::{Peer, RequestContext, Service, ServiceRole};
|
||||
|
||||
pub struct TowerHandler<S, R: ServiceRole> {
|
||||
pub service: S,
|
||||
pub info: R::Info,
|
||||
pub peer: Option<Peer<R>>,
|
||||
role: PhantomData<R>,
|
||||
}
|
||||
|
||||
impl<S, R: ServiceRole> TowerHandler<S, R> {
|
||||
pub fn new(service: S, info: R::Info) -> Self {
|
||||
Self {
|
||||
service,
|
||||
role: PhantomData,
|
||||
info,
|
||||
peer: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, R: ServiceRole> Service<R> for TowerHandler<S, R>
|
||||
where
|
||||
S: TowerService<R::PeerReq, Response = R::Resp> + Sync + Send + Clone + 'static,
|
||||
S::Error: Into<crate::Error>,
|
||||
S::Future: Send,
|
||||
{
|
||||
async fn handle_request(
|
||||
&self,
|
||||
request: R::PeerReq,
|
||||
_context: RequestContext<R>,
|
||||
) -> Result<R::Resp, crate::Error> {
|
||||
let mut service = self.service.clone();
|
||||
poll_fn(|cx| service.poll_ready(cx))
|
||||
.await
|
||||
.map_err(Into::into)?;
|
||||
let resp = service.call(request).await.map_err(Into::into)?;
|
||||
Ok(resp)
|
||||
}
|
||||
|
||||
fn handle_notification(
|
||||
&self,
|
||||
_notification: R::PeerNot,
|
||||
) -> impl Future<Output = Result<(), crate::Error>> + Send + '_ {
|
||||
std::future::ready(Ok(()))
|
||||
}
|
||||
|
||||
fn get_peer(&self) -> Option<Peer<R>> {
|
||||
self.peer.clone()
|
||||
}
|
||||
|
||||
fn set_peer(&mut self, peer: Peer<R>) {
|
||||
self.peer = Some(peer);
|
||||
}
|
||||
|
||||
fn get_info(&self) -> R::Info {
|
||||
self.info.clone()
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue