openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-01 19:33:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use HttpClient directly for OAuth bootstrap

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Ahmed Ibrahim
2026-04-22 19:18:23 -07:00
co-authored by Codex
parent b8780250d6
commit f8ebe83fc9
64 changed files with 21967 additions and 314 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
-2
View File
@@ -10559,8 +10559,6 @@ dependencies = [
[[package]]
name = "rmcp"
version = "0.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1bef41ebc9ebed2c1b1d90203e9d1756091e8a00bbc3107676151f39868ca0ee"
dependencies = [
"async-trait",
"axum",
codex-rs/Cargo.toml
+1 -2
View File
@@ -478,8 +478,7 @@ ratatui = { git = "https://github.com/nornagon/ratatui", rev = "9b2ad1298408c459
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "132f5b39c862e3a970f731d709608b3e6276d5f6" }
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "9200079d3b54a1ff51072e24d81fd354f085156f" }
# Uncomment to debug local changes.
# rmcp = { path = "../../rust-sdk/crates/rmcp" }
rmcp = { path = "vendor/rmcp" }
[patch."ssh://git@github.com/openai-oss-forks/tungstenite-rs.git"]
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "9200079d3b54a1ff51072e24d81fd354f085156f" }
codex-rs/rmcp-client/src/oauth.rs
-4
View File
@@ -48,7 +48,6 @@ use codex_keyring_store::KeyringStore;
use rmcp::transport::auth::AuthorizationManager;
use tokio::sync::Mutex;
use crate::oauth_http_client::OAuthHttpProxy;
use codex_utils_home_dir::find_codex_home;
const KEYRING_SERVICE: &str = "Codex MCP Credentials";
@@ -259,7 +258,6 @@ struct OAuthPersistorInner {
server_name: String,
url: String,
authorization_manager: Arc<Mutex<AuthorizationManager>>,
_oauth_http_proxy: Option<OAuthHttpProxy>,
store_mode: OAuthCredentialsStoreMode,
last_credentials: Mutex<Option<StoredOAuthTokens>>,
}
@@ -269,7 +267,6 @@ impl OAuthPersistor {
server_name: String,
url: String,
authorization_manager: Arc<Mutex<AuthorizationManager>>,
oauth_http_proxy: Option<OAuthHttpProxy>,
store_mode: OAuthCredentialsStoreMode,
initial_credentials: Option<StoredOAuthTokens>,
) -> Self {
@@ -278,7 +275,6 @@ impl OAuthPersistor {
server_name,
url,
authorization_manager,
_oauth_http_proxy: oauth_http_proxy,
store_mode,
last_credentials: Mutex::new(initial_credentials),
}),
codex-rs/rmcp-client/src/oauth_http_client.rs
+99 -301
View File
@@ -1,278 +1,121 @@
//! OAuth bootstrap helpers that route non-browser HTTP through the shared
//! `HttpClient` capability.
//! OAuth bootstrap helpers built on the shared `HttpClient` capability.
//!
//! The browser-facing part of MCP OAuth stays on the orchestrator, but the
//! HTTP requests used for discovery, registration, token exchange, and refresh
//! need to follow the selected MCP placement. This module discovers OAuth
//! metadata through `HttpClient` and exposes token/registration endpoints
//! through a small localhost proxy so RMCP's OAuth manager can keep using its
//! existing `reqwest`-based internals.
//! Browser launch, callback handling, and token persistence remain
//! orchestrator-owned. This module covers the non-browser OAuth HTTP work
//! needed for discovery, registration, token exchange, and refresh so those
//! requests follow the same local-or-remote placement as the MCP transport.
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::thread;
use anyhow::Result;
use anyhow::anyhow;
use async_trait::async_trait;
use codex_exec_server::HttpClient;
use codex_exec_server::HttpHeader;
use codex_exec_server::HttpRequestParams;
use oauth2::HttpRequest;
use oauth2::HttpResponse;
use oauth2::http::HeaderName as OAuthHeaderName;
use oauth2::http::HeaderValue as OAuthHeaderValue;
use reqwest::Method;
use reqwest::StatusCode;
use reqwest::Url;
use reqwest::header::HeaderMap;
use rmcp::transport::auth::AuthError;
use rmcp::transport::auth::AuthorizationManager;
use rmcp::transport::auth::AuthorizationMetadata;
use rmcp::transport::auth::CredentialStore;
use rmcp::transport::auth::OAuthHttpClient;
use rmcp::transport::auth::StoredCredentials;
use tiny_http::Header;
use tiny_http::Response;
use tiny_http::Server;
const DISCOVERY_TIMEOUT_MS: u64 = 30_000;
const MCP_PROTOCOL_VERSION_HEADER: &str = "MCP-Protocol-Version";
const MCP_PROTOCOL_VERSION: &str = "2024-11-05";
const TOKEN_PROXY_PATH: &str = "/oauth/token";
const REGISTRATION_PROXY_PATH: &str = "/oauth/register";
const DEFAULT_TIMEOUT_MS: u64 = 30_000;
pub(crate) struct OAuthHttpSetup {
pub(crate) authorization_manager: AuthorizationManager,
pub(crate) proxy: Option<OAuthHttpProxy>,
}
pub(crate) async fn create_oauth_http_setup(
pub(crate) async fn create_oauth_authorization_manager(
url: &str,
default_headers: &HeaderMap,
default_headers: HeaderMap,
http_client: Arc<dyn HttpClient>,
initial_credentials: Option<StoredCredentials>,
) -> Result<OAuthHttpSetup> {
let metadata = discover_authorization_metadata(url, default_headers, Arc::clone(&http_client))
.await
.map_err(|error| anyhow!(error))?;
let proxy = OAuthHttpProxy::new(Arc::clone(&http_client), default_headers, &metadata)?;
) -> Result<AuthorizationManager> {
let mut authorization_manager = AuthorizationManager::new(url)
.await
.map_err(|error| anyhow!(error))?;
authorization_manager.set_metadata(proxy.rewrite_metadata(&metadata)?);
authorization_manager
.with_oauth_http_client(HttpClientOAuthAdapter::new(http_client, default_headers))
.map_err(|error| anyhow!(error))?;
if let Some(initial_credentials) = initial_credentials {
let client_id = initial_credentials.client_id.clone();
authorization_manager.set_credential_store(StaticCredentialStore::new(initial_credentials));
authorization_manager
.configure_client_id(&client_id)
.initialize_from_store()
.await
.map_err(|error| anyhow!(error))?;
}
Ok(OAuthHttpSetup {
authorization_manager,
proxy: Some(proxy),
})
Ok(authorization_manager)
}
async fn discover_authorization_metadata(
url: &str,
default_headers: &HeaderMap,
#[derive(Clone)]
struct HttpClientOAuthAdapter {
http_client: Arc<dyn HttpClient>,
) -> Result<AuthorizationMetadata, AuthError> {
let base_url = Url::parse(url).map_err(|error| AuthError::OAuthError(error.to_string()))?;
default_headers: HeaderMap,
}
for discovery_url in generate_discovery_urls(&base_url) {
let response = http_client
impl HttpClientOAuthAdapter {
fn new(http_client: Arc<dyn HttpClient>, default_headers: HeaderMap) -> Self {
Self {
http_client,
default_headers,
}
}
}
#[async_trait]
impl OAuthHttpClient for HttpClientOAuthAdapter {
async fn execute(&self, request: HttpRequest) -> Result<HttpResponse, AuthError> {
let method = Method::from_bytes(request.method().as_str().as_bytes())
.map_err(|error| AuthError::InternalError(error.to_string()))?;
let url = Url::parse(&request.uri().to_string()).map_err(AuthError::UrlError)?;
let mut headers = protocol_headers(&self.default_headers);
for (name, value) in request.headers() {
headers.push(HttpHeader {
name: name.as_str().to_string(),
value: value
.to_str()
.map_err(|error| AuthError::InternalError(error.to_string()))?
.to_string(),
});
}
let response = self
.http_client
.http_request(HttpRequestParams {
method: Method::GET.as_str().to_string(),
url: discovery_url.to_string(),
headers: discovery_headers(default_headers)?,
body: None,
timeout_ms: Some(DISCOVERY_TIMEOUT_MS),
request_id: next_request_id("oauth-discovery"),
method: method.as_str().to_string(),
url: url.to_string(),
headers,
body: Some(request.body().clone().into()),
timeout_ms: Some(DEFAULT_TIMEOUT_MS),
request_id: next_request_id("oauth-http"),
stream_response: false,
})
.await
.map_err(|error| AuthError::OAuthError(error.to_string()))?;
if response.status != StatusCode::OK.as_u16() {
continue;
let mut http_response = HttpResponse::new(response.body.into_inner());
*http_response.status_mut() = reqwest::StatusCode::from_u16(response.status)
.map_err(|error| AuthError::InternalError(error.to_string()))?;
for header in response.headers {
let header_name = header
.name
.parse::<OAuthHeaderName>()
.map_err(|error| AuthError::InternalError(error.to_string()))?;
let header_value = header
.value
.parse::<OAuthHeaderValue>()
.map_err(|error| AuthError::InternalError(error.to_string()))?;
http_response
.headers_mut()
.insert(header_name, header_value);
}
let metadata = serde_json::from_slice::<AuthorizationMetadata>(&response.body.into_inner())
.map_err(|error| AuthError::OAuthError(error.to_string()))?;
return Ok(metadata);
}
Err(AuthError::NoAuthorizationSupport)
}
fn generate_discovery_urls(base_url: &Url) -> Vec<Url> {
let mut candidates = Vec::new();
let trimmed = base_url
.path()
.trim_start_matches('/')
.trim_end_matches('/');
let mut push_candidate = |discovery_path: String| {
let mut discovery_url = base_url.clone();
discovery_url.set_query(None);
discovery_url.set_fragment(None);
discovery_url.set_path(&discovery_path);
candidates.push(discovery_url);
};
if trimmed.is_empty() {
push_candidate("/.well-known/oauth-authorization-server".to_string());
push_candidate("/.well-known/openid-configuration".to_string());
} else {
push_candidate(format!("/.well-known/oauth-authorization-server/{trimmed}"));
push_candidate(format!("/.well-known/openid-configuration/{trimmed}"));
push_candidate(format!("/{trimmed}/.well-known/openid-configuration"));
push_candidate("/.well-known/oauth-authorization-server".to_string());
}
candidates
}
fn discovery_headers(default_headers: &HeaderMap) -> Result<Vec<HttpHeader>, AuthError> {
let mut headers = protocol_headers(default_headers);
headers.push(HttpHeader {
name: MCP_PROTOCOL_VERSION_HEADER.to_string(),
value: MCP_PROTOCOL_VERSION.to_string(),
});
Ok(headers)
}
pub(crate) struct OAuthHttpProxy {
server: Arc<Server>,
base_url: String,
routes: Arc<HashMap<String, String>>,
}
impl OAuthHttpProxy {
fn new(
http_client: Arc<dyn HttpClient>,
default_headers: &HeaderMap,
metadata: &AuthorizationMetadata,
) -> Result<Self> {
let server = Arc::new(Server::http("127.0.0.1:0").map_err(|error| anyhow!(error))?);
let address = match server.server_addr() {
tiny_http::ListenAddr::IP(address) => address,
#[cfg(not(target_os = "windows"))]
_ => return Err(anyhow!("unable to determine OAuth HTTP proxy bind address")),
};
let base_url = format!("http://{address}");
let mut routes = HashMap::new();
routes.insert(
TOKEN_PROXY_PATH.to_string(),
metadata.token_endpoint.clone(),
);
if let Some(registration_endpoint) = metadata.registration_endpoint.clone() {
routes.insert(REGISTRATION_PROXY_PATH.to_string(), registration_endpoint);
}
let routes = Arc::new(routes);
let server_for_thread = Arc::clone(&server);
let routes_for_thread = Arc::clone(&routes);
let default_headers = default_headers.clone();
let runtime = tokio::runtime::Handle::current();
thread::spawn(move || {
while let Ok(mut request) = server_for_thread.recv() {
let request_url = request.url().to_string();
let route_key = request_url
.split_once('?')
.map(|(path, _)| path)
.unwrap_or(request_url.as_str())
.to_string();
let Some(target_url) = routes_for_thread.get(&route_key).cloned() else {
let _ = request.respond(Response::empty(StatusCode::NOT_FOUND.as_u16()));
continue;
};
let mut target = target_url;
if let Some((_, query)) = request_url.split_once('?') {
target.push('?');
target.push_str(query);
}
let method = request.method().as_str().to_string();
let mut body = Vec::new();
if request.as_reader().read_to_end(&mut body).is_err() {
let _ = request.respond(Response::empty(StatusCode::BAD_REQUEST.as_u16()));
continue;
}
let mut headers = protocol_headers(&default_headers);
for header in request.headers() {
let field = header.field.as_str().to_string();
if field.eq_ignore_ascii_case("host")
|| field.eq_ignore_ascii_case("content-length")
{
continue;
}
headers.push(HttpHeader {
name: field,
value: header.value.to_string(),
});
}
let response = runtime.block_on(http_client.http_request(HttpRequestParams {
method,
url: target,
headers,
body: Some(body.into()),
timeout_ms: Some(DISCOVERY_TIMEOUT_MS),
request_id: next_request_id("oauth-proxy"),
stream_response: false,
}));
let response = match response {
Ok(response) => response,
Err(_) => {
let _ = request.respond(Response::empty(StatusCode::BAD_GATEWAY.as_u16()));
continue;
}
};
let mut proxy_response = Response::from_data(response.body.into_inner())
.with_status_code(response.status);
for header in response.headers {
if let Ok(header) = Header::from_bytes(header.name.as_bytes(), header.value) {
proxy_response.add_header(header);
}
}
let _ = request.respond(proxy_response);
}
});
Ok(Self {
server,
base_url,
routes,
})
}
fn rewrite_metadata(&self, metadata: &AuthorizationMetadata) -> Result<AuthorizationMetadata> {
let mut rewritten = metadata.clone();
rewritten.token_endpoint = self.proxied_url(TOKEN_PROXY_PATH)?;
rewritten.registration_endpoint = if self.routes.contains_key(REGISTRATION_PROXY_PATH) {
Some(self.proxied_url(REGISTRATION_PROXY_PATH)?)
} else {
None
};
Ok(rewritten)
}
fn proxied_url(&self, path: &str) -> Result<String> {
let mut url = Url::parse(&self.base_url).map_err(|error| anyhow!(error))?;
url.set_path(path);
Ok(url.to_string())
}
}
impl Drop for OAuthHttpProxy {
fn drop(&mut self) {
self.server.unblock();
Ok(http_response)
}
}
@@ -289,8 +132,8 @@ impl StaticCredentialStore {
}
}
#[async_trait::async_trait]
impl CredentialStore for StaticCredentialStore {
#[async_trait]
impl rmcp::transport::auth::CredentialStore for StaticCredentialStore {
async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
Ok(self.stored.lock().await.clone())
}
@@ -333,6 +176,9 @@ mod tests {
use codex_exec_server::HttpRequestResponse;
use futures::FutureExt;
use futures::future::BoxFuture;
use oauth2::http::Method as OAuthMethod;
use oauth2::http::Uri;
use oauth2::http::header::CONTENT_TYPE;
use pretty_assertions::assert_eq;
use super::*;
@@ -385,84 +231,36 @@ mod tests {
}
#[tokio::test]
async fn discover_authorization_metadata_uses_http_client() -> Result<()> {
async fn oauth_http_adapter_uses_http_client() -> Result<()> {
let http_client = TestHttpClient::default();
http_client.push_response(Ok(HttpRequestResponse {
status: StatusCode::OK.as_u16(),
headers: Vec::new(),
body: serde_json::to_vec(&AuthorizationMetadata {
authorization_endpoint: "https://example.com/authorize".to_string(),
token_endpoint: "https://example.com/token".to_string(),
registration_endpoint: None,
issuer: None,
jwks_uri: None,
scopes_supported: None,
response_types_supported: None,
additional_fields: HashMap::new(),
})?
.into(),
}));
let metadata = discover_authorization_metadata(
"https://example.com/mcp",
&HeaderMap::new(),
Arc::new(http_client.clone()),
)
.await?;
assert_eq!(
metadata.authorization_endpoint,
"https://example.com/authorize"
);
assert_eq!(metadata.token_endpoint, "https://example.com/token");
assert_eq!(
http_client.requests()[0].url,
"https://example.com/.well-known/oauth-authorization-server/mcp"
);
Ok(())
}
#[tokio::test]
async fn oauth_http_proxy_forwards_token_requests() -> Result<()> {
let http_client = TestHttpClient::default();
http_client.push_response(Ok(HttpRequestResponse {
status: StatusCode::OK.as_u16(),
status: reqwest::StatusCode::OK.as_u16(),
headers: vec![HttpHeader {
name: "content-type".to_string(),
value: "application/json".to_string(),
}],
body: br#"{"access_token":"abc","token_type":"Bearer"}"#.to_vec().into(),
body: br#"{"ok":true}"#.to_vec().into(),
}));
let proxy = OAuthHttpProxy::new(
Arc::new(http_client.clone()),
&HeaderMap::new(),
&AuthorizationMetadata {
authorization_endpoint: "https://example.com/authorize".to_string(),
token_endpoint: "https://remote.example.com/token".to_string(),
registration_endpoint: None,
issuer: None,
jwks_uri: None,
scopes_supported: None,
response_types_supported: None,
additional_fields: HashMap::new(),
},
)?;
let adapter = HttpClientOAuthAdapter::new(Arc::new(http_client.clone()), HeaderMap::new());
let mut request = HttpRequest::new(b"grant_type=refresh_token".to_vec());
*request.method_mut() = OAuthMethod::POST;
*request.uri_mut() = "https://example.com/token"
.parse::<Uri>()
.map_err(|error| anyhow!(error))?;
request.headers_mut().insert(
CONTENT_TYPE,
OAuthHeaderValue::from_static("application/x-www-form-urlencoded"),
);
let response = reqwest::Client::new()
.post(proxy.proxied_url(TOKEN_PROXY_PATH)?)
.header("content-type", "application/x-www-form-urlencoded")
.body("grant_type=refresh_token")
.send()
.await?;
let response = adapter.execute(request).await?;
assert_eq!(response.status(), StatusCode::OK);
let requests = http_client.requests();
assert_eq!(requests.len(), 1);
assert_eq!(requests[0].url, "https://remote.example.com/token");
assert_eq!(response.status(), reqwest::StatusCode::OK);
assert_eq!(response.body().as_slice(), br#"{"ok":true}"#);
assert_eq!(http_client.requests().len(), 1);
assert_eq!(http_client.requests()[0].url, "https://example.com/token");
assert_eq!(
String::from_utf8(requests[0].body.clone().unwrap().into_inner())?,
String::from_utf8(http_client.requests()[0].body.clone().unwrap().into_inner())?,
"grant_type=refresh_token"
);
Ok(())
codex-rs/rmcp-client/src/rmcp_client.rs
+3 -5
View File
@@ -63,7 +63,7 @@ use crate::http_client_adapter::StreamableHttpClientAdapterError;
use crate::load_oauth_tokens;
use crate::oauth::OAuthPersistor;
use crate::oauth::StoredOAuthTokens;
use crate::oauth_http_client::create_oauth_http_setup;
use crate::oauth_http_client::create_oauth_authorization_manager;
use crate::stdio_server_launcher::StdioServerCommand;
use crate::stdio_server_launcher::StdioServerLauncher;
use crate::stdio_server_launcher::StdioServerTransport;
@@ -933,9 +933,9 @@ async fn create_oauth_transport_and_runtime(
StreamableHttpClientTransport<AuthClient<StreamableHttpClientAdapter>>,
OAuthPersistor,
)> {
let oauth_http_setup = create_oauth_http_setup(
let manager = create_oauth_authorization_manager(
url,
&default_headers,
default_headers.clone(),
http_client.clone(),
Some(StoredCredentials {
client_id: initial_tokens.client_id.clone(),
@@ -943,7 +943,6 @@ async fn create_oauth_transport_and_runtime(
}),
)
.await?;
let manager = oauth_http_setup.authorization_manager;
let auth_client = AuthClient::new(
StreamableHttpClientAdapter::new(http_client, default_headers),
@@ -960,7 +959,6 @@ async fn create_oauth_transport_and_runtime(
server_name.to_string(),
url.to_string(),
auth_manager,
oauth_http_setup.proxy,
credentials_store,
Some(initial_tokens),
);
codex-rs/vendor/rmcp/Cargo.toml
Vendored
+486
View File
@@ -0,0 +1,486 @@
# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO
#
# When uploading crates to the registry Cargo will automatically
# "normalize" Cargo.toml files for maximal compatibility
# with all versions of Cargo and also rewrite `path` dependencies
# to registry (e.g., crates.io) dependencies.
#
# If you are reading this file be aware that the original Cargo.toml
# will likely look very different (and much more reasonable).
# See Cargo.toml.orig for the original contents.
[package]
edition = "2024"
name = "rmcp"
version = "0.15.0"
build = "build.rs"
autolib = false
autobins = false
autoexamples = false
autotests = false
autobenches = false
description = "Rust SDK for Model Context Protocol"
homepage = "https://github.com/modelcontextprotocol/rust-sdk"
documentation = "https://docs.rs/rmcp"
readme = "README.md"
license = "Apache-2.0"
repository = "https://github.com/modelcontextprotocol/rust-sdk/"
resolver = "2"
[package.metadata.docs.rs]
all-features = true
rustdoc-args = [
"--cfg",
"docsrs",
]
[features]
__reqwest = ["dep:reqwest"]
auth = [
"dep:oauth2",
"__reqwest",
"dep:http",
"dep:url",
]
client = ["dep:tokio-stream"]
client-side-sse = [
"dep:sse-stream",
"dep:http",
]
default = [
"base64",
"macros",
"server",
]
elicitation = ["dep:url"]
macros = [
"dep:rmcp-macros",
"dep:pastey",
]
reqwest = [
"__reqwest",
"reqwest?/rustls-tls",
]
reqwest-native-tls = [
"__reqwest",
"reqwest?/native-tls",
]
reqwest-tls-no-provider = [
"__reqwest",
"reqwest?/rustls-tls-no-provider",
]
schemars = ["dep:schemars"]
server = [
"transport-async-rw",
"dep:schemars",
"dep:pastey",
]
server-side-http = [
"uuid",
"dep:rand",
"dep:tokio-stream",
"dep:http",
"dep:http-body",
"dep:http-body-util",
"dep:bytes",
"dep:sse-stream",
"dep:axum",
"tower",
]
tower = ["dep:tower-service"]
transport-async-rw = [
"tokio/io-util",
"tokio-util/codec",
]
transport-child-process = [
"transport-async-rw",
"tokio/process",
"dep:process-wrap",
]
transport-io = [
"transport-async-rw",
"tokio/io-std",
]
transport-streamable-http-client = [
"client-side-sse",
"transport-worker",
]
transport-streamable-http-client-reqwest = [
"transport-streamable-http-client",
"__reqwest",
]
transport-streamable-http-server = [
"transport-streamable-http-server-session",
"server-side-http",
"transport-worker",
]
transport-streamable-http-server-session = [
"transport-async-rw",
"dep:tokio-stream",
]
transport-worker = ["dep:tokio-stream"]
[lib]
name = "rmcp"
path = "src/lib.rs"
[[test]]
name = "test_client_initialization"
path = "tests/test_client_initialization.rs"
[[test]]
name = "test_close_connection"
path = "tests/test_close_connection.rs"
required-features = [
"server",
"client",
]
[[test]]
name = "test_completion"
path = "tests/test_completion.rs"
[[test]]
name = "test_complex_schema"
path = "tests/test_complex_schema.rs"
[[test]]
name = "test_custom_request"
path = "tests/test_custom_request.rs"
required-features = [
"server",
"client",
]
[[test]]
name = "test_deserialization"
path = "tests/test_deserialization.rs"
[[test]]
name = "test_elicitation"
path = "tests/test_elicitation.rs"
required-features = [
"elicitation",
"client",
"server",
]
[[test]]
name = "test_embedded_resource_meta"
path = "tests/test_embedded_resource_meta.rs"
[[test]]
name = "test_handler_wrappers"
path = "tests/test_handler_wrappers.rs"
[[test]]
name = "test_json_schema_detection"
path = "tests/test_json_schema_detection.rs"
[[test]]
name = "test_logging"
path = "tests/test_logging.rs"
required-features = [
"server",
"client",
]
[[test]]
name = "test_message_protocol"
path = "tests/test_message_protocol.rs"
required-features = ["client"]
[[test]]
name = "test_message_schema"
path = "tests/test_message_schema.rs"
required-features = [
"server",
"client",
"schemars",
]
[[test]]
name = "test_notification"
path = "tests/test_notification.rs"
required-features = [
"server",
"client",
]
[[test]]
name = "test_progress_subscriber"
path = "tests/test_progress_subscriber.rs"
required-features = [
"server",
"client",
"macros",
]
[[test]]
name = "test_prompt_handler"
path = "tests/test_prompt_handler.rs"
[[test]]
name = "test_prompt_macro_annotations"
path = "tests/test_prompt_macro_annotations.rs"
[[test]]
name = "test_prompt_macros"
path = "tests/test_prompt_macros.rs"
required-features = [
"server",
"client",
]
[[test]]
name = "test_prompt_routers"
path = "tests/test_prompt_routers.rs"
[[test]]
name = "test_resource_link"
path = "tests/test_resource_link.rs"
[[test]]
name = "test_resource_link_integration"
path = "tests/test_resource_link_integration.rs"
[[test]]
name = "test_sampling"
path = "tests/test_sampling.rs"
required-features = [
"server",
"client",
]
[[test]]
name = "test_streamable_http_priming"
path = "tests/test_streamable_http_priming.rs"
required-features = [
"server",
"client",
"transport-streamable-http-server",
"reqwest",
]
[[test]]
name = "test_structured_output"
path = "tests/test_structured_output.rs"
[[test]]
name = "test_task"
path = "tests/test_task.rs"
required-features = [
"server",
"client",
"macros",
]
[[test]]
name = "test_task_support_validation"
path = "tests/test_task_support_validation.rs"
[[test]]
name = "test_tool_builder_methods"
path = "tests/test_tool_builder_methods.rs"
[[test]]
name = "test_tool_handler"
path = "tests/test_tool_handler.rs"
[[test]]
name = "test_tool_macro_annotations"
path = "tests/test_tool_macro_annotations.rs"
[[test]]
name = "test_tool_macros"
path = "tests/test_tool_macros.rs"
required-features = [
"server",
"client",
]
[[test]]
name = "test_tool_result_meta"
path = "tests/test_tool_result_meta.rs"
[[test]]
name = "test_tool_routers"
path = "tests/test_tool_routers.rs"
[[test]]
name = "test_with_js"
path = "tests/test_with_js.rs"
required-features = [
"server",
"client",
"transport-child-process",
"transport-streamable-http-server",
"transport-streamable-http-client",
"__reqwest",
]
[[test]]
name = "test_with_python"
path = "tests/test_with_python.rs"
required-features = [
"reqwest",
"server",
"client",
"transport-child-process",
]
[dependencies.async-trait]
version = "0.1.89"
[dependencies.axum]
version = "0.8"
features = []
optional = true
[dependencies.base64]
version = "0.22"
optional = true
[dependencies.bytes]
version = "1"
optional = true
[dependencies.futures]
version = "0.3"
[dependencies.http]
version = "1"
optional = true
[dependencies.http-body]
version = "1"
optional = true
[dependencies.http-body-util]
version = "0.1"
optional = true
[dependencies.oauth2]
version = "5.0"
features = ["reqwest"]
optional = true
default-features = false
[dependencies.pastey]
version = "0.2.0"
optional = true
[dependencies.pin-project-lite]
version = "0.2"
[dependencies.process-wrap]
version = "9.0"
features = ["tokio1"]
optional = true
[dependencies.rand]
version = "0.9"
optional = true
[dependencies.reqwest]
version = "0.12"
features = [
"json",
"stream",
]
optional = true
default-features = false
[dependencies.rmcp-macros]
version = "0.15.0"
optional = true
[dependencies.schemars]
version = "1.0"
features = ["chrono04"]
optional = true
[dependencies.serde]
version = "1.0"
features = [
"derive",
"rc",
]
[dependencies.serde_json]
version = "1.0"
[dependencies.sse-stream]
version = "0.2"
optional = true
[dependencies.thiserror]
version = "2"
[dependencies.tokio]
version = "1"
features = [
"sync",
"macros",
"rt",
"time",
]
[dependencies.tokio-stream]
version = "0.1"
optional = true
[dependencies.tokio-util]
version = "0.7"
[dependencies.tower-service]
version = "0.3"
optional = true
[dependencies.tracing]
version = "0.1"
[dependencies.url]
version = "2.4"
optional = true
[dependencies.uuid]
version = "1"
features = ["v4"]
optional = true
[dev-dependencies.anyhow]
version = "1.0"
[dev-dependencies.async-trait]
version = "0.1"
[dev-dependencies.schemars]
version = "1.1.0"
features = ["chrono04"]
[dev-dependencies.tokio]
version = "1"
features = ["full"]
[dev-dependencies.tracing-subscriber]
version = "0.3"
features = [
"env-filter",
"std",
"fmt",
]
[target.'cfg(all(target_family = "wasm", target_os = "unknown"))'.dependencies.chrono]
version = "0.4.38"
features = [
"serde",
"clock",
"std",
"oldtime",
]
default-features = false
[target.'cfg(not(all(target_family = "wasm", target_os = "unknown")))'.dependencies.chrono]
version = "0.4.38"
features = ["serde"]
codex-rs/vendor/rmcp/README.md
Vendored
+310
View File
@@ -0,0 +1,310 @@
<style>
.rustdoc-hidden { display: none; }
</style>
<div class="rustdoc-hidden">
# RMCP: Rust Model Context Protocol
[![Crates.io](https://img.shields.io/crates/v/rmcp.svg)](https://crates.io/crates/rmcp)
[![Documentation](https://docs.rs/rmcp/badge.svg)](https://docs.rs/rmcp)
</div>
`rmcp` is the official Rust implementation of the Model Context Protocol (MCP), a protocol designed for AI assistants to communicate with other services. This library can be used to build both servers that expose capabilities to AI assistants and clients that interact with such servers.
## Quick Start
### Server Implementation
Creating a server with tools is simple using the `#[tool]` macro:
```rust,no_run
use rmcp::{
ServerHandler, ServiceExt,
handler::server::tool::ToolRouter,
model::*,
tool, tool_handler, tool_router,
transport::stdio,
ErrorData as McpError,
};
use std::sync::Arc;
use tokio::sync::Mutex;
#[derive(Clone)]
pub struct Counter {
counter: Arc<Mutex<i32>>,
tool_router: ToolRouter<Self>,
}
#[tool_router]
impl Counter {
fn new() -> Self {
Self {
counter: Arc::new(Mutex::new(0)),
tool_router: Self::tool_router(),
}
}
#[tool(description = "Increment the counter by 1")]
async fn increment(&self) -> Result<CallToolResult, McpError> {
let mut counter = self.counter.lock().await;
*counter += 1;
Ok(CallToolResult::success(vec![Content::text(
counter.to_string(),
)]))
}
#[tool(description = "Get the current counter value")]
async fn get(&self) -> Result<CallToolResult, McpError> {
let counter = self.counter.lock().await;
Ok(CallToolResult::success(vec![Content::text(
counter.to_string(),
)]))
}
}
// Implement the server handler
#[tool_handler]
impl ServerHandler for Counter {
fn get_info(&self) -> ServerInfo {
ServerInfo {
instructions: Some("A simple counter that tallies the number of times the increment tool has been used".into()),
capabilities: ServerCapabilities::builder().enable_tools().build(),
..Default::default()
}
}
}
// Run the server
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Create and run the server with STDIO transport
let service = Counter::new().serve(stdio()).await.inspect_err(|e| {
println!("Error starting server: {}", e);
})?;
service.waiting().await?;
Ok(())
}
```
### Structured Output
Tools can return structured JSON data with schemas. Use the [`Json`] wrapper:
```rust
# use rmcp::{tool, tool_router, handler::server::{tool::ToolRouter, wrapper::Parameters}, Json};
# use schemars::JsonSchema;
# use serde::{Serialize, Deserialize};
#
#[derive(Serialize, Deserialize, JsonSchema)]
struct CalculationRequest {
a: i32,
b: i32,
operation: String,
}
#[derive(Serialize, Deserialize, JsonSchema)]
struct CalculationResult {
result: i32,
operation: String,
}
# #[derive(Clone)]
# struct Calculator {
# tool_router: ToolRouter<Self>,
# }
#
# #[tool_router]
# impl Calculator {
#[tool(name = "calculate", description = "Perform a calculation")]
async fn calculate(&self, params: Parameters<CalculationRequest>) -> Result<Json<CalculationResult>, String> {
let result = match params.0.operation.as_str() {
"add" => params.0.a + params.0.b,
"multiply" => params.0.a * params.0.b,
_ => return Err("Unknown operation".to_string()),
};
Ok(Json(CalculationResult { result, operation: params.0.operation }))
}
# }
```
The `#[tool]` macro automatically generates an output schema from the `CalculationResult` type.
## Tasks
RMCP implements the task lifecycle from SEP-1686 so long-running or asynchronous tool calls can be queued and polled safely.
- **Create:** set the `task` field on `CallToolRequestParam` to ask the server to enqueue the tool call. The response is a `CreateTaskResult` that includes the generated `task.task_id`.
- **Inspect:** use `tasks/get` (`GetTaskInfoRequest`) to retrieve metadata such as status, timestamps, TTL, and poll interval.
- **Await results:** call `tasks/result` (`GetTaskResultRequest`) to block until the task completes and receive either the final `CallToolResult` payload or a protocol error.
- **Cancel:** call `tasks/cancel` (`CancelTaskRequest`) to request termination of a running task.
To expose task support, enable the `tasks` capability when building `ServerCapabilities`. The `#[task_handler]` macro and `OperationProcessor` utility provide reference implementations for enqueuing, tracking, and collecting task results.
### Client Implementation
Creating a client to interact with a server:
```rust,no_run
use rmcp::{
ServiceExt,
model::CallToolRequestParams,
transport::{ConfigureCommandExt, TokioChildProcess},
};
use tokio::process::Command;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Connect to a server running as a child process
let service = ()
.serve(TokioChildProcess::new(Command::new("uvx").configure(
|cmd| {
cmd.arg("mcp-server-git");
},
))?)
.await?;
// Get server information
let server_info = service.peer_info();
println!("Connected to server: {server_info:#?}");
// List available tools
let tools = service.list_tools(Default::default()).await?;
println!("Available tools: {tools:#?}");
// Call a tool
let result = service
.call_tool(CallToolRequestParams {
meta: None,
name: "git_status".into(),
arguments: serde_json::json!({ "repo_path": "." }).as_object().cloned(),
task: None,
})
.await?;
println!("Result: {result:#?}");
// Gracefully close the connection
service.cancel().await?;
Ok(())
}
```
For more examples, see the [examples directory](https://github.com/anthropics/mcp-rust-sdk/tree/main/examples) in the repository.
## Transport Options
RMCP supports multiple transport mechanisms, each suited for different use cases:
### `transport-async-rw`
Low-level interface for asynchronous read/write operations. This is the foundation for many other transports.
### `transport-io`
For working directly with I/O streams (`tokio::io::AsyncRead` and `tokio::io::AsyncWrite`).
### `transport-child-process`
Run MCP servers as child processes and communicate via standard I/O.
Example:
```rust,ignore
use rmcp::transport::TokioChildProcess;
use tokio::process::Command;
let transport = TokioChildProcess::new(Command::new("mcp-server"))?;
let service = client.serve(transport).await?;
```
## Access with peer interface when handling message
You can get the [`Peer`](crate::service::Peer) struct from [`NotificationContext`](crate::service::NotificationContext) and [`RequestContext`](crate::service::RequestContext).
```rust, ignore
# use rmcp::{
# ServerHandler,
# model::{LoggingLevel, LoggingMessageNotificationParam, ProgressNotificationParam},
# service::{NotificationContext, RoleServer},
# };
# pub struct Handler;
impl ServerHandler for Handler {
async fn on_progress(
&self,
notification: ProgressNotificationParam,
context: NotificationContext<RoleServer>,
) {
let peer = context.peer;
let _ = peer
.notify_logging_message(LoggingMessageNotificationParam {
level: LoggingLevel::Info,
logger: None,
data: serde_json::json!({
"message": format!("Progress: {}", notification.progress),
}),
})
.await;
}
}
```
## 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();
```
## Feature Flags
RMCP uses feature flags to control which components are included:
- `client`: Enable client functionality
- `server`: Enable server functionality and the tool system
- `macros`: Enable the `#[tool]` macro (enabled by default)
- Transport-specific features:
- `transport-async-rw`: Async read/write support
- `transport-io`: I/O stream support
- `transport-child-process`: Child process support
- `transport-streamable-http-client` / `transport-streamable-http-server`: HTTP streaming (client agnostic, see [`StreamableHttpClientTransport`](crate::transport::StreamableHttpClientTransport) for details)
- `transport-streamable-http-client-reqwest`: a default `reqwest` implementation of the streamable http client
- `auth`: OAuth2 authentication support
- `schemars`: JSON Schema generation (for tool definitions)
- TLS backend options (for HTTP transports):
- `reqwest`: Uses rustls (pure Rust TLS, recommended default)
- `reqwest-native-tls`: Uses platform native TLS (OpenSSL on Linux, Secure Transport on macOS, SChannel on Windows)
- `reqwest-tls-no-provider`: Uses rustls without a default crypto provider (bring your own)
## Transports
- `transport-io`: Server stdio transport
- `transport-child-process`: Client stdio transport
- `transport-streamable-http-server` streamable http server transport
- `transport-streamable-http-client` streamable http client transport
<details>
<summary>Transport</summary>
The transport type must implement the [`Transport`](crate::transport::Transport) trait, which allows it to send messages concurrently and receive messages sequentially.
There are 2 pairs of standard transport types:
| transport | client | server |
|:---------------:|:-----------------------------------------------------------------------------------:|:-----------------------------------------------------------------------------:|
| std IO | [`TokioChildProcess`](crate::transport::TokioChildProcess) | [`stdio`](crate::transport::stdio) |
| streamable http | [`StreamableHttpClientTransport`](crate::transport::StreamableHttpClientTransport) | [`StreamableHttpService`](crate::transport::StreamableHttpService) |
#### [`IntoTransport`](crate::transport::IntoTransport) trait
[`IntoTransport`](crate::transport::IntoTransport) is a helper trait that implicitly converts a type into a transport type.
These types automatically implement [`IntoTransport`](crate::transport::IntoTransport):
1. A type that implements both `futures::Sink` and `futures::Stream`, or a tuple `(Tx, Rx)` where `Tx` is `futures::Sink` and `Rx` is `futures::Stream`.
2. A type that implements both `tokio::io::AsyncRead` and `tokio::io::AsyncWrite`, or a tuple `(R, W)` where `R` is `tokio::io::AsyncRead` and `W` is `tokio::io::AsyncWrite`.
3. A type that implements the [`Worker`](crate::transport::worker::Worker) trait.
4. A type that implements the [`Transport`](crate::transport::Transport) trait.
</details>
## License
This project is licensed under the terms specified in the repository's LICENSE file.
codex-rs/vendor/rmcp/build.rs
Vendored
+23
View File
@@ -0,0 +1,23 @@
// Install git hooks on build
fn main() {
// Only run in the workspace root (not when building as a dependency)
let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
let workspace_root = std::path::Path::new(&manifest_dir)
.parent()
.and_then(|p| p.parent());
if let Some(root) = workspace_root {
let githooks_dir = root.join(".githooks");
let git_dir = root.join(".git");
// Only configure if we're in the actual workspace (not a dependency)
// and git directory exists
if githooks_dir.exists() && git_dir.exists() {
// Configure git to use our hooks directory
let _ = std::process::Command::new("git")
.args(["config", "core.hooksPath", ".githooks"])
.current_dir(root)
.output();
}
}
}
codex-rs/vendor/rmcp/src/error.rs
Vendored
+61
View File
@@ -0,0 +1,61 @@
use std::{borrow::Cow, fmt::Display};
pub use crate::model::ErrorData;
#[deprecated(
note = "Use `rmcp::ErrorData` instead, `rmcp::ErrorData` could become `RmcpError` in the future."
)]
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 {}
/// This is an unified error type for the errors could be returned by the service.
#[derive(Debug, thiserror::Error)]
#[allow(clippy::large_enum_variant)]
pub enum RmcpError {
#[cfg(any(feature = "client", feature = "server"))]
#[error("Service error: {0}")]
Service(#[from] crate::ServiceError),
#[cfg(feature = "client")]
#[error("Client initialization error: {0}")]
ClientInitialize(#[from] crate::service::ClientInitializeError),
#[cfg(feature = "server")]
#[error("Server initialization error: {0}")]
ServerInitialize(#[from] crate::service::ServerInitializeError),
#[error("Runtime error: {0}")]
Runtime(#[from] tokio::task::JoinError),
#[error("Transport creation error: {error}")]
// TODO: Maybe we can introduce something like `TryIntoTransport` to auto wrap transport type,
// but it could be an breaking change, so we could do it in the future.
TransportCreation {
into_transport_type_name: Cow<'static, str>,
into_transport_type_id: std::any::TypeId,
#[source]
error: Box<dyn std::error::Error + Send + Sync>,
},
// and cancellation shouldn't be an error?
// TODO: add more error variants as needed
#[error("Task error: {0}")]
TaskError(String),
}
impl RmcpError {
pub fn transport_creation<T: 'static>(
error: impl Into<Box<dyn std::error::Error + Send + Sync>>,
) -> Self {
RmcpError::TransportCreation {
into_transport_type_id: std::any::TypeId::of::<T>(),
into_transport_type_name: std::any::type_name::<T>().into(),
error: error.into(),
}
}
}
codex-rs/vendor/rmcp/src/handler.rs
Vendored
+6
View File
@@ -0,0 +1,6 @@
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub mod client;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub mod server;
codex-rs/vendor/rmcp/src/handler/client.rs
Vendored
+376
View File
@@ -0,0 +1,376 @@
pub mod progress;
use std::sync::Arc;
use crate::{
error::ErrorData as McpError,
model::*,
service::{NotificationContext, 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(Box::new)
.map(ClientResult::CreateMessageResult),
ServerRequest::ListRootsRequest(_) => self
.list_roots(context)
.await
.map(ClientResult::ListRootsResult),
ServerRequest::CreateElicitationRequest(request) => self
.create_elicitation(request.params, context)
.await
.map(ClientResult::CreateElicitationResult),
ServerRequest::CustomRequest(request) => self
.on_custom_request(request, context)
.await
.map(ClientResult::CustomResult),
}
}
async fn handle_notification(
&self,
notification: <RoleClient as ServiceRole>::PeerNot,
context: NotificationContext<RoleClient>,
) -> Result<(), McpError> {
match notification {
ServerNotification::CancelledNotification(notification) => {
self.on_cancelled(notification.params, context).await
}
ServerNotification::ProgressNotification(notification) => {
self.on_progress(notification.params, context).await
}
ServerNotification::LoggingMessageNotification(notification) => {
self.on_logging_message(notification.params, context).await
}
ServerNotification::ResourceUpdatedNotification(notification) => {
self.on_resource_updated(notification.params, context).await
}
ServerNotification::ResourceListChangedNotification(_notification_no_param) => {
self.on_resource_list_changed(context).await
}
ServerNotification::ToolListChangedNotification(_notification_no_param) => {
self.on_tool_list_changed(context).await
}
ServerNotification::PromptListChangedNotification(_notification_no_param) => {
self.on_prompt_list_changed(context).await
}
ServerNotification::ElicitationCompletionNotification(notification) => {
self.on_url_elicitation_notification_complete(notification.params, context)
.await
}
ServerNotification::CustomNotification(notification) => {
self.on_custom_notification(notification, context).await
}
};
Ok(())
}
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: CreateMessageRequestParams,
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()))
}
/// Handle an elicitation request from a server asking for user input.
///
/// This method is called when a server needs interactive input from the user
/// during tool execution. Implementations should present the message to the user,
/// collect their input according to the requested schema, and return the result.
///
/// # Arguments
/// * `request` - The elicitation request with message and schema
/// * `context` - The request context
///
/// # Returns
/// The user's response including action (accept/decline/cancel) and optional data
///
/// # Default Behavior
/// The default implementation automatically declines all elicitation requests.
/// Real clients should override this to provide user interaction.
///
/// # Example
/// ```rust,ignore
/// use rmcp::model::CreateElicitationRequestParam;
/// use rmcp::{
/// model::ErrorData as McpError,
/// model::*,
/// service::{NotificationContext, RequestContext, RoleClient, Service, ServiceRole},
/// };
/// use rmcp::ClientHandler;
///
/// impl ClientHandler for MyClient {
/// async fn create_elicitation(
/// &self,
/// request: CreateElicitationRequestParam,
/// context: RequestContext<RoleClient>,
/// ) -> Result<CreateElicitationResult, McpError> {
/// match request {
/// CreateElicitationRequestParam::FormElicitationParam {meta, message, requested_schema,} => {
/// // Display message to user and collect input according to requested_schema
/// let user_input = get_user_input(message, requested_schema).await?;
/// Ok(CreateElicitationResult {
/// action: ElicitationAction::Accept,
/// content: Some(user_input),
/// })
/// }
/// CreateElicitationRequestParam::UrlElicitationParam {meta, message, url, elicitation_id,} => {
/// // Open URL in browser for user to complete elicitation
/// open_url_in_browser(url).await?;
/// Ok(CreateElicitationResult {
/// action: ElicitationAction::Accept,
/// content: None,
/// })
/// }
/// }
/// }
/// }
/// ```
fn create_elicitation(
&self,
request: CreateElicitationRequestParams,
context: RequestContext<RoleClient>,
) -> impl Future<Output = Result<CreateElicitationResult, McpError>> + Send + '_ {
// Default implementation declines all requests - real clients should override this
let _ = (request, context);
std::future::ready(Ok(CreateElicitationResult {
action: ElicitationAction::Decline,
content: None,
}))
}
fn on_custom_request(
&self,
request: CustomRequest,
context: RequestContext<RoleClient>,
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
let CustomRequest { method, .. } = request;
let _ = context;
std::future::ready(Err(McpError::new(
ErrorCode::METHOD_NOT_FOUND,
method,
None,
)))
}
fn on_cancelled(
&self,
params: CancelledNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_progress(
&self,
params: ProgressNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_logging_message(
&self,
params: LoggingMessageNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_resource_updated(
&self,
params: ResourceUpdatedNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_resource_list_changed(
&self,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_tool_list_changed(
&self,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_prompt_list_changed(
&self,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_url_elicitation_notification_complete(
&self,
params: ElicitationResponseNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_custom_notification(
&self,
notification: CustomNotification,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
let _ = (notification, context);
std::future::ready(())
}
fn get_info(&self) -> ClientInfo {
ClientInfo::default()
}
}
/// Do nothing, with default client info.
impl ClientHandler for () {}
/// Do nothing, with a specific client info.
impl ClientHandler for ClientInfo {
fn get_info(&self) -> ClientInfo {
self.clone()
}
}
macro_rules! impl_client_handler_for_wrapper {
($wrapper:ident) => {
impl<T: ClientHandler> ClientHandler for $wrapper<T> {
fn ping(
&self,
context: RequestContext<RoleClient>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
(**self).ping(context)
}
fn create_message(
&self,
params: CreateMessageRequestParams,
context: RequestContext<RoleClient>,
) -> impl Future<Output = Result<CreateMessageResult, McpError>> + Send + '_ {
(**self).create_message(params, context)
}
fn list_roots(
&self,
context: RequestContext<RoleClient>,
) -> impl Future<Output = Result<ListRootsResult, McpError>> + Send + '_ {
(**self).list_roots(context)
}
fn create_elicitation(
&self,
request: CreateElicitationRequestParams,
context: RequestContext<RoleClient>,
) -> impl Future<Output = Result<CreateElicitationResult, McpError>> + Send + '_ {
(**self).create_elicitation(request, context)
}
fn on_custom_request(
&self,
request: CustomRequest,
context: RequestContext<RoleClient>,
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
(**self).on_custom_request(request, context)
}
fn on_cancelled(
&self,
params: CancelledNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_cancelled(params, context)
}
fn on_progress(
&self,
params: ProgressNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_progress(params, context)
}
fn on_logging_message(
&self,
params: LoggingMessageNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_logging_message(params, context)
}
fn on_resource_updated(
&self,
params: ResourceUpdatedNotificationParam,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_resource_updated(params, context)
}
fn on_resource_list_changed(
&self,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_resource_list_changed(context)
}
fn on_tool_list_changed(
&self,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_tool_list_changed(context)
}
fn on_prompt_list_changed(
&self,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_prompt_list_changed(context)
}
fn on_custom_notification(
&self,
notification: CustomNotification,
context: NotificationContext<RoleClient>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_custom_notification(notification, context)
}
fn get_info(&self) -> ClientInfo {
(**self).get_info()
}
}
};
}
impl_client_handler_for_wrapper!(Box);
impl_client_handler_for_wrapper!(Arc);
codex-rs/vendor/rmcp/src/handler/client/progress.rs
Vendored
+100
View File
@@ -0,0 +1,100 @@
use std::{collections::HashMap, sync::Arc};
use futures::{Stream, StreamExt};
use tokio::sync::RwLock;
use tokio_stream::wrappers::ReceiverStream;
use crate::model::{ProgressNotificationParam, ProgressToken};
type Dispatcher =
Arc<RwLock<HashMap<ProgressToken, tokio::sync::mpsc::Sender<ProgressNotificationParam>>>>;
/// A dispatcher for progress notifications.
#[derive(Debug, Clone, Default)]
pub struct ProgressDispatcher {
pub(crate) dispatcher: Dispatcher,
}
impl ProgressDispatcher {
const CHANNEL_SIZE: usize = 16;
pub fn new() -> Self {
Self::default()
}
/// Handle a progress notification by sending it to the appropriate subscriber
pub async fn handle_notification(&self, notification: ProgressNotificationParam) {
let token = &notification.progress_token;
if let Some(sender) = self.dispatcher.read().await.get(token).cloned() {
let send_result = sender.send(notification).await;
if let Err(e) = send_result {
tracing::warn!("Failed to send progress notification: {e}");
}
}
}
/// Subscribe to progress notifications for a specific token.
///
/// If you drop the returned `ProgressSubscriber`, it will automatically unsubscribe from notifications for that token.
pub async fn subscribe(&self, progress_token: ProgressToken) -> ProgressSubscriber {
let (sender, receiver) = tokio::sync::mpsc::channel(Self::CHANNEL_SIZE);
self.dispatcher
.write()
.await
.insert(progress_token.clone(), sender);
let receiver = ReceiverStream::new(receiver);
ProgressSubscriber {
progress_token,
receiver,
dispatcher: self.dispatcher.clone(),
}
}
/// Unsubscribe from progress notifications for a specific token.
pub async fn unsubscribe(&self, token: &ProgressToken) {
self.dispatcher.write().await.remove(token);
}
/// Clear all dispatcher.
pub async fn clear(&self) {
let mut dispatcher = self.dispatcher.write().await;
dispatcher.clear();
}
}
pub struct ProgressSubscriber {
pub(crate) progress_token: ProgressToken,
pub(crate) receiver: ReceiverStream<ProgressNotificationParam>,
pub(crate) dispatcher: Dispatcher,
}
impl ProgressSubscriber {
pub fn progress_token(&self) -> &ProgressToken {
&self.progress_token
}
}
impl Stream for ProgressSubscriber {
type Item = ProgressNotificationParam;
fn poll_next(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
self.receiver.poll_next_unpin(cx)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.receiver.size_hint()
}
}
impl Drop for ProgressSubscriber {
fn drop(&mut self) {
let token = self.progress_token.clone();
self.receiver.close();
let dispatcher = self.dispatcher.clone();
tokio::spawn(async move {
let mut dispatcher = dispatcher.write_owned().await;
dispatcher.remove(&token);
});
}
}
codex-rs/vendor/rmcp/src/handler/server.rs
Vendored
+570
View File
@@ -0,0 +1,570 @@
use std::sync::Arc;
use crate::{
error::ErrorData as McpError,
model::{TaskSupport, *},
service::{NotificationContext, RequestContext, RoleServer, Service, ServiceRole},
};
pub mod common;
pub mod prompt;
mod resource;
pub mod router;
pub mod tool;
pub mod tool_name_validation;
pub mod wrapper;
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) => {
let is_task = request.params.task.is_some();
// Validate task support mode per MCP specification
if let Some(tool) = self.get_tool(&request.params.name) {
match (tool.task_support(), is_task) {
// If taskSupport is "required", clients MUST invoke the tool as a task.
// Servers MUST return a -32601 (Method not found) error if they don't.
(TaskSupport::Required, false) => {
return Err(McpError::new(
ErrorCode::METHOD_NOT_FOUND,
"Tool requires task-based invocation",
None,
));
}
// If taskSupport is "forbidden" (default), clients MUST NOT invoke as a task.
(TaskSupport::Forbidden, true) => {
return Err(McpError::invalid_params(
"Tool does not support task-based invocation",
None,
));
}
_ => {}
}
}
if is_task {
tracing::info!("Enqueueing task for tool call: {}", request.params.name);
self.enqueue_task(request.params, context.clone())
.await
.map(ServerResult::CreateTaskResult)
} else {
self.call_tool(request.params, context)
.await
.map(ServerResult::CallToolResult)
}
}
ClientRequest::ListToolsRequest(request) => self
.list_tools(request.params, context)
.await
.map(ServerResult::ListToolsResult),
ClientRequest::CustomRequest(request) => self
.on_custom_request(request, context)
.await
.map(ServerResult::CustomResult),
ClientRequest::ListTasksRequest(request) => self
.list_tasks(request.params, context)
.await
.map(ServerResult::ListTasksResult),
ClientRequest::GetTaskInfoRequest(request) => self
.get_task_info(request.params, context)
.await
.map(ServerResult::GetTaskInfoResult),
ClientRequest::GetTaskResultRequest(request) => self
.get_task_result(request.params, context)
.await
.map(ServerResult::TaskResult),
ClientRequest::CancelTaskRequest(request) => self
.cancel_task(request.params, context)
.await
.map(ServerResult::empty),
}
}
async fn handle_notification(
&self,
notification: <RoleServer as ServiceRole>::PeerNot,
context: NotificationContext<RoleServer>,
) -> Result<(), McpError> {
match notification {
ClientNotification::CancelledNotification(notification) => {
self.on_cancelled(notification.params, context).await
}
ClientNotification::ProgressNotification(notification) => {
self.on_progress(notification.params, context).await
}
ClientNotification::InitializedNotification(_notification) => {
self.on_initialized(context).await
}
ClientNotification::RootsListChangedNotification(_notification) => {
self.on_roots_list_changed(context).await
}
ClientNotification::CustomNotification(notification) => {
self.on_custom_notification(notification, context).await
}
};
Ok(())
}
fn get_info(&self) -> <RoleServer as ServiceRole>::Info {
self.get_info()
}
}
#[allow(unused_variables)]
pub trait ServerHandler: Sized + Send + Sync + 'static {
fn enqueue_task(
&self,
_request: CallToolRequestParams,
_context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CreateTaskResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::internal_error(
"Task processing not implemented".to_string(),
None,
)))
}
fn ping(
&self,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
std::future::ready(Ok(()))
}
// handle requests
fn initialize(
&self,
request: InitializeRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<InitializeResult, McpError>> + Send + '_ {
if context.peer.peer_info().is_none() {
context.peer.set_peer_info(request);
}
std::future::ready(Ok(self.get_info()))
}
fn complete(
&self,
request: CompleteRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CompleteResult, McpError>> + Send + '_ {
std::future::ready(Ok(CompleteResult::default()))
}
fn set_level(
&self,
request: SetLevelRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<SetLevelRequestMethod>()))
}
fn get_prompt(
&self,
request: GetPromptRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetPromptResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<GetPromptRequestMethod>()))
}
fn list_prompts(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListPromptsResult, McpError>> + Send + '_ {
std::future::ready(Ok(ListPromptsResult::default()))
}
fn list_resources(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListResourcesResult, McpError>> + Send + '_ {
std::future::ready(Ok(ListResourcesResult::default()))
}
fn list_resource_templates(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListResourceTemplatesResult, McpError>> + Send + '_ {
std::future::ready(Ok(ListResourceTemplatesResult::default()))
}
fn read_resource(
&self,
request: ReadResourceRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ReadResourceResult, McpError>> + Send + '_ {
std::future::ready(Err(
McpError::method_not_found::<ReadResourceRequestMethod>(),
))
}
fn subscribe(
&self,
request: SubscribeRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<SubscribeRequestMethod>()))
}
fn unsubscribe(
&self,
request: UnsubscribeRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<UnsubscribeRequestMethod>()))
}
fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CallToolResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<CallToolRequestMethod>()))
}
fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
std::future::ready(Ok(ListToolsResult::default()))
}
/// Get a tool definition by name.
///
/// The default implementation returns `None`, which bypasses validation.
/// When using `#[tool_handler]`, this method is automatically implemented.
fn get_tool(&self, _name: &str) -> Option<Tool> {
None
}
fn on_custom_request(
&self,
request: CustomRequest,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
let CustomRequest { method, .. } = request;
let _ = context;
std::future::ready(Err(McpError::new(
ErrorCode::METHOD_NOT_FOUND,
method,
None,
)))
}
fn on_cancelled(
&self,
notification: CancelledNotificationParam,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_progress(
&self,
notification: ProgressNotificationParam,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_initialized(
&self,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
tracing::info!("client initialized");
std::future::ready(())
}
fn on_roots_list_changed(
&self,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
std::future::ready(())
}
fn on_custom_notification(
&self,
notification: CustomNotification,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
let _ = (notification, context);
std::future::ready(())
}
fn get_info(&self) -> ServerInfo {
ServerInfo::default()
}
fn list_tasks(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListTasksResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<ListTasksMethod>()))
}
fn get_task_info(
&self,
request: GetTaskInfoParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetTaskInfoResult, McpError>> + Send + '_ {
std::future::ready(Err(McpError::method_not_found::<GetTaskInfoMethod>()))
}
fn get_task_result(
&self,
request: GetTaskResultParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<TaskResult, McpError>> + Send + '_ {
let _ = (request, context);
std::future::ready(Err(McpError::method_not_found::<GetTaskResultMethod>()))
}
fn cancel_task(
&self,
request: CancelTaskParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
let _ = (request, context);
std::future::ready(Err(McpError::method_not_found::<CancelTaskMethod>()))
}
}
macro_rules! impl_server_handler_for_wrapper {
($wrapper:ident) => {
impl<T: ServerHandler> ServerHandler for $wrapper<T> {
fn enqueue_task(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CreateTaskResult, McpError>> + Send + '_ {
(**self).enqueue_task(request, context)
}
fn ping(
&self,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
(**self).ping(context)
}
fn initialize(
&self,
request: InitializeRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<InitializeResult, McpError>> + Send + '_ {
(**self).initialize(request, context)
}
fn complete(
&self,
request: CompleteRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CompleteResult, McpError>> + Send + '_ {
(**self).complete(request, context)
}
fn set_level(
&self,
request: SetLevelRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
(**self).set_level(request, context)
}
fn get_prompt(
&self,
request: GetPromptRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetPromptResult, McpError>> + Send + '_ {
(**self).get_prompt(request, context)
}
fn list_prompts(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListPromptsResult, McpError>> + Send + '_ {
(**self).list_prompts(request, context)
}
fn list_resources(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListResourcesResult, McpError>> + Send + '_ {
(**self).list_resources(request, context)
}
fn list_resource_templates(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListResourceTemplatesResult, McpError>> + Send + '_
{
(**self).list_resource_templates(request, context)
}
fn read_resource(
&self,
request: ReadResourceRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ReadResourceResult, McpError>> + Send + '_ {
(**self).read_resource(request, context)
}
fn subscribe(
&self,
request: SubscribeRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
(**self).subscribe(request, context)
}
fn unsubscribe(
&self,
request: UnsubscribeRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
(**self).unsubscribe(request, context)
}
fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CallToolResult, McpError>> + Send + '_ {
(**self).call_tool(request, context)
}
fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
(**self).list_tools(request, context)
}
fn get_tool(&self, name: &str) -> Option<Tool> {
(**self).get_tool(name)
}
fn on_custom_request(
&self,
request: CustomRequest,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<CustomResult, McpError>> + Send + '_ {
(**self).on_custom_request(request, context)
}
fn on_cancelled(
&self,
notification: CancelledNotificationParam,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_cancelled(notification, context)
}
fn on_progress(
&self,
notification: ProgressNotificationParam,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_progress(notification, context)
}
fn on_initialized(
&self,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_initialized(context)
}
fn on_roots_list_changed(
&self,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_roots_list_changed(context)
}
fn on_custom_notification(
&self,
notification: CustomNotification,
context: NotificationContext<RoleServer>,
) -> impl Future<Output = ()> + Send + '_ {
(**self).on_custom_notification(notification, context)
}
fn get_info(&self) -> ServerInfo {
(**self).get_info()
}
fn list_tasks(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListTasksResult, McpError>> + Send + '_ {
(**self).list_tasks(request, context)
}
fn get_task_info(
&self,
request: GetTaskInfoParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<GetTaskInfoResult, McpError>> + Send + '_ {
(**self).get_task_info(request, context)
}
fn get_task_result(
&self,
request: GetTaskResultParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<TaskResult, McpError>> + Send + '_ {
(**self).get_task_result(request, context)
}
fn cancel_task(
&self,
request: CancelTaskParams,
context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
(**self).cancel_task(request, context)
}
}
};
}
impl_server_handler_for_wrapper!(Box);
impl_server_handler_for_wrapper!(Arc);
codex-rs/vendor/rmcp/src/handler/server/common.rs
Vendored
+271
View File
@@ -0,0 +1,271 @@
//! Common utilities shared between tool and prompt handlers
use std::{any::TypeId, collections::HashMap, sync::Arc};
use schemars::JsonSchema;
use crate::{
RoleServer, model::JsonObject, schemars::generate::SchemaSettings, service::RequestContext,
};
/// Generates a JSON schema for a type
pub fn 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 {
// explicitly to align json schema version to official specifications.
// refer to https://github.com/modelcontextprotocol/modelcontextprotocol/pull/655 for details.
let mut settings = SchemaSettings::draft2020_12();
settings.transforms = vec![Box::new(schemars::transform::AddNullable::default())];
let generator = settings.into_generator();
let schema = generator.into_root_schema_for::<T>();
let object = serde_json::to_value(schema).expect("failed to serialize schema");
let object = match object {
serde_json::Value::Object(object) => object,
_ => panic!(
"Schema serialization produced non-object value: expected JSON object but got {:?}",
object
),
};
let schema = Arc::new(object);
cache
.write()
.expect("schema cache lock poisoned")
.insert(TypeId::of::<T>(), schema.clone());
schema
}
})
}
/// Generate and validate a JSON schema for outputSchema (must have root type "object").
pub fn schema_for_output<T: JsonSchema + std::any::Any>() -> Result<Arc<JsonObject>, String> {
thread_local! {
static CACHE_FOR_OUTPUT: std::sync::RwLock<HashMap<TypeId, Result<Arc<JsonObject>, String>>> = Default::default();
};
CACHE_FOR_OUTPUT.with(|cache| {
// Try to get from cache first
if let Some(result) = cache
.read()
.expect("output schema cache lock poisoned")
.get(&TypeId::of::<T>())
{
return result.clone();
}
// Generate and validate schema
let schema = schema_for_type::<T>();
let result = match schema.get("type") {
Some(serde_json::Value::String(t)) if t == "object" => Ok(schema.clone()),
Some(serde_json::Value::String(t)) => Err(format!(
"MCP specification requires tool outputSchema to have root type 'object', but found '{}'.",
t
)),
None => Err(
"Schema is missing 'type' field. MCP specification requires outputSchema to have root type 'object'.".to_string()
),
Some(other) => Err(format!(
"Schema 'type' field has unexpected format: {:?}. Expected \"object\".",
other
)),
};
// Cache the result (both success and error cases)
cache
.write()
.expect("output schema cache lock poisoned")
.insert(TypeId::of::<T>(), result.clone());
result
})
}
/// Trait for extracting parts from a context, unifying tool and prompt extraction
pub trait FromContextPart<C>: Sized {
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData>;
}
/// Common extractors that can be used by both tool and prompt handlers
impl<C> FromContextPart<C> for RequestContext<RoleServer>
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().clone())
}
}
impl<C> FromContextPart<C> for tokio_util::sync::CancellationToken
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().ct.clone())
}
}
impl<C> FromContextPart<C> for crate::model::Extensions
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().extensions.clone())
}
}
pub struct Extension<T>(pub T);
impl<C, T> FromContextPart<C> for Extension<T>
where
C: AsRequestContext,
T: Send + Sync + 'static + Clone,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
let extension = context
.as_request_context()
.extensions
.get::<T>()
.cloned()
.ok_or_else(|| {
crate::ErrorData::invalid_params(
format!("missing extension {}", std::any::type_name::<T>()),
None,
)
})?;
Ok(Extension(extension))
}
}
impl<C> FromContextPart<C> for crate::Peer<RoleServer>
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(context.as_request_context().peer.clone())
}
}
impl<C> FromContextPart<C> for crate::model::Meta
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
let request_context = context.as_request_context_mut();
let mut meta = crate::model::Meta::default();
std::mem::swap(&mut meta, &mut request_context.meta);
Ok(meta)
}
}
pub struct RequestId(pub crate::model::RequestId);
impl<C> FromContextPart<C> for RequestId
where
C: AsRequestContext,
{
fn from_context_part(context: &mut C) -> Result<Self, crate::ErrorData> {
Ok(RequestId(context.as_request_context().id.clone()))
}
}
/// Trait for types that can provide access to RequestContext
pub trait AsRequestContext {
fn as_request_context(&self) -> &RequestContext<RoleServer>;
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer>;
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(serde::Serialize, serde::Deserialize, JsonSchema)]
struct TestObject {
value: i32,
}
#[derive(serde::Serialize, serde::Deserialize, JsonSchema)]
struct AnotherTestObject {
value: i32,
}
#[test]
fn test_schema_for_type_handles_primitive() {
let schema = schema_for_type::<i32>();
assert_eq!(schema.get("type"), Some(&serde_json::json!("integer")));
}
#[test]
fn test_schema_for_type_handles_array() {
let schema = schema_for_type::<Vec<i32>>();
assert_eq!(schema.get("type"), Some(&serde_json::json!("array")));
let items = schema.get("items").and_then(|v| v.as_object());
assert_eq!(
items.unwrap().get("type"),
Some(&serde_json::json!("integer"))
);
}
#[test]
fn test_schema_for_type_handles_struct() {
let schema = schema_for_type::<TestObject>();
assert_eq!(schema.get("type"), Some(&serde_json::json!("object")));
let properties = schema.get("properties").and_then(|v| v.as_object());
assert!(properties.unwrap().contains_key("value"));
}
#[test]
fn test_schema_for_type_caches_primitive_types() {
let schema1 = schema_for_type::<i32>();
let schema2 = schema_for_type::<i32>();
assert!(Arc::ptr_eq(&schema1, &schema2));
}
#[test]
fn test_schema_for_type_caches_struct_types() {
let schema1 = schema_for_type::<TestObject>();
let schema2 = schema_for_type::<TestObject>();
assert!(Arc::ptr_eq(&schema1, &schema2));
}
#[test]
fn test_schema_for_type_different_types_different_schemas() {
let schema1 = schema_for_type::<TestObject>();
let schema2 = schema_for_type::<AnotherTestObject>();
assert!(!Arc::ptr_eq(&schema1, &schema2));
}
#[test]
fn test_schema_for_type_arc_can_be_shared() {
let schema = schema_for_type::<TestObject>();
let cloned = schema.clone();
assert!(Arc::ptr_eq(&schema, &cloned));
}
#[test]
fn test_schema_for_output_rejects_primitive() {
let result = schema_for_output::<i32>();
assert!(result.is_err(),);
}
#[test]
fn test_schema_for_output_accepts_object() {
let result = schema_for_output::<TestObject>();
assert!(result.is_ok(),);
}
}
codex-rs/vendor/rmcp/src/handler/server/prompt.rs
Vendored
+367
View File
@@ -0,0 +1,367 @@
//! Prompt handling infrastructure for MCP servers
//!
//! This module provides the core types and traits for implementing prompt handlers
//! in MCP servers. Prompts allow servers to provide reusable templates for LLM
//! interactions with customizable arguments.
use std::{future::Future, marker::PhantomData};
use futures::future::{BoxFuture, FutureExt};
use serde::de::DeserializeOwned;
use super::common::{AsRequestContext, FromContextPart};
pub use super::common::{Extension, RequestId};
use crate::{
RoleServer,
handler::server::wrapper::Parameters,
model::{GetPromptResult, PromptMessage},
service::RequestContext,
};
/// Context for prompt retrieval operations
pub struct PromptContext<'a, S> {
pub server: &'a S,
pub name: String,
pub arguments: Option<serde_json::Map<String, serde_json::Value>>,
pub context: RequestContext<RoleServer>,
}
impl<'a, S> PromptContext<'a, S> {
pub fn new(
server: &'a S,
name: String,
arguments: Option<serde_json::Map<String, serde_json::Value>>,
context: RequestContext<RoleServer>,
) -> Self {
Self {
server,
name,
arguments,
context,
}
}
}
impl<S> AsRequestContext for PromptContext<'_, S> {
fn as_request_context(&self) -> &RequestContext<RoleServer> {
&self.context
}
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer> {
&mut self.context
}
}
/// Trait for handling prompt retrieval
pub trait GetPromptHandler<S, A> {
fn handle(
self,
context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>;
}
/// Type alias for dynamic prompt handlers
pub type DynGetPromptHandler<S> = dyn for<'a> Fn(PromptContext<'a, S>) -> BoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
+ Send
+ Sync;
/// Adapter type for async methods that return `Vec<PromptMessage>`
pub struct AsyncMethodAdapter<T>(PhantomData<T>);
/// Adapter type for async methods with parameters that return `Vec<PromptMessage>`
pub struct AsyncMethodWithArgsAdapter<T>(PhantomData<T>);
/// Adapter types for macro-generated implementations
#[allow(clippy::type_complexity)]
pub struct AsyncPromptAdapter<P, Fut, R>(PhantomData<fn(P) -> fn(Fut) -> R>);
pub struct SyncPromptAdapter<P, R>(PhantomData<fn(P) -> R>);
pub struct AsyncPromptMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
pub struct SyncPromptMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
/// Trait for types that can be converted into GetPromptResult
pub trait IntoGetPromptResult {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData>;
}
impl IntoGetPromptResult for GetPromptResult {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
Ok(self)
}
}
impl IntoGetPromptResult for Vec<PromptMessage> {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
Ok(GetPromptResult {
description: None,
messages: self,
})
}
}
impl<T: IntoGetPromptResult> IntoGetPromptResult for Result<T, crate::ErrorData> {
fn into_get_prompt_result(self) -> Result<GetPromptResult, crate::ErrorData> {
self.and_then(|v| v.into_get_prompt_result())
}
}
// Future wrapper that automatically handles IntoGetPromptResult conversion
pin_project_lite::pin_project! {
#[project = IntoGetPromptResultFutProj]
pub enum IntoGetPromptResultFut<F, R> {
Pending {
#[pin]
fut: F,
_marker: PhantomData<R>,
},
Ready {
#[pin]
result: futures::future::Ready<Result<GetPromptResult, crate::ErrorData>>,
}
}
}
impl<F, R> Future for IntoGetPromptResultFut<F, R>
where
F: Future<Output = R>,
R: IntoGetPromptResult,
{
type Output = Result<GetPromptResult, crate::ErrorData>;
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
match self.project() {
IntoGetPromptResultFutProj::Pending { fut, _marker } => fut
.poll(cx)
.map(IntoGetPromptResult::into_get_prompt_result),
IntoGetPromptResultFutProj::Ready { result } => result.poll(cx),
}
}
}
// Prompt-specific extractor for prompt name
pub struct PromptName(pub String);
impl<S> FromContextPart<PromptContext<'_, S>> for PromptName {
fn from_context_part(context: &mut PromptContext<S>) -> Result<Self, crate::ErrorData> {
Ok(Self(context.name.clone()))
}
}
// Special implementation for Parameters that handles prompt arguments
impl<S, P> FromContextPart<PromptContext<'_, S>> for Parameters<P>
where
P: DeserializeOwned,
{
fn from_context_part(context: &mut PromptContext<S>) -> Result<Self, crate::ErrorData> {
let params = if let Some(args_map) = context.arguments.take() {
let args_value = serde_json::Value::Object(args_map);
serde_json::from_value::<P>(args_value).map_err(|e| {
crate::ErrorData::invalid_params(format!("Failed to parse parameters: {}", e), None)
})?
} else {
// Try to deserialize from empty object for optional fields
serde_json::from_value::<P>(serde_json::json!({})).map_err(|e| {
crate::ErrorData::invalid_params(
format!("Missing required parameters: {}", e),
None,
)
})?
};
Ok(Parameters(params))
}
}
// Macro to generate GetPromptHandler implementations for various parameter combinations
macro_rules! impl_prompt_handler_for {
($($T: ident)*) => {
impl_prompt_handler_for!([] [$($T)*]);
};
// finished
([$($Tn: ident)*] []) => {
impl_prompt_handler_for!(@impl $($Tn)*);
};
([$($Tn: ident)*] [$Tn_1: ident $($Rest: ident)*]) => {
impl_prompt_handler_for!(@impl $($Tn)*);
impl_prompt_handler_for!([$($Tn)* $Tn_1] [$($Rest)*]);
};
(@impl $($Tn: ident)*) => {
// Implementation for async methods (transformed by #[prompt] macro)
impl<$($Tn,)* S, F, R> GetPromptHandler<S, ($($Tn,)*)> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send,
)*
F: FnOnce(&S, $($Tn,)*) -> BoxFuture<'_, R> + Send,
R: IntoGetPromptResult + Send + 'static,
S: Send + Sync + 'static,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let service = context.server;
let fut = self(service, $($Tn,)*);
async move {
let result = fut.await;
result.into_get_prompt_result()
}.boxed()
}
}
// Implementation for sync methods
impl<$($Tn,)* S, F, R> GetPromptHandler<S, SyncPromptMethodAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send,
)*
F: FnOnce(&S, $($Tn,)*) -> R + Send,
R: IntoGetPromptResult + Send,
S: Send + Sync,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let service = context.server;
let result = self(service, $($Tn,)*);
std::future::ready(result.into_get_prompt_result()).boxed()
}
}
// AsyncPromptAdapter - for standalone functions returning GetPromptResult
impl<$($Tn,)* S, F, Fut, R> GetPromptHandler<S, AsyncPromptAdapter<($($Tn,)*), Fut, R>> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send + 'static,
)*
F: FnOnce($($Tn,)*) -> Fut + Send + 'static,
Fut: Future<Output = Result<R, crate::ErrorData>> + Send + 'static,
R: IntoGetPromptResult + Send + 'static,
S: Send + Sync + 'static,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
// Extract all parameters before moving into the async block
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
// Since we're dealing with standalone functions that don't take &S,
// we can return a 'static future
Box::pin(async move {
let result = self($($Tn,)*).await?;
result.into_get_prompt_result()
})
}
}
// SyncPromptAdapter - for standalone sync functions returning Result
impl<$($Tn,)* S, F, R> GetPromptHandler<S, SyncPromptAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: for<'a> FromContextPart<PromptContext<'a, S>> + Send + 'static,
)*
F: FnOnce($($Tn,)*) -> Result<R, crate::ErrorData> + Send + 'static,
R: IntoGetPromptResult + Send + 'static,
S: Send + Sync,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn handle(
self,
mut context: PromptContext<'_, S>,
) -> BoxFuture<'_, Result<GetPromptResult, crate::ErrorData>>
{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let result = self($($Tn,)*);
std::future::ready(result.and_then(|r| r.into_get_prompt_result())).boxed()
}
}
};
}
// Invoke the macro to generate implementations for up to 16 parameters
impl_prompt_handler_for!(T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 T12 T13 T14 T15);
/// Extract prompt arguments from a type's JSON schema
/// This function analyzes the schema of a type and extracts the properties
/// as PromptArgument entries with name, description, and required status
pub fn cached_arguments_from_schema<T: schemars::JsonSchema + std::any::Any>()
-> Option<Vec<crate::model::PromptArgument>> {
let schema = super::common::schema_for_type::<T>();
let schema_value = serde_json::Value::Object((*schema).clone());
let properties = schema_value.get("properties").and_then(|p| p.as_object());
if let Some(props) = properties {
let required = schema_value
.get("required")
.and_then(|r| r.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str())
.collect::<std::collections::HashSet<_>>()
})
.unwrap_or_default();
let mut arguments = Vec::new();
for (name, prop_schema) in props {
let description = prop_schema
.get("description")
.and_then(|d| d.as_str())
.map(|s| s.to_string());
arguments.push(crate::model::PromptArgument {
name: name.clone(),
title: None,
description,
required: Some(required.contains(name.as_str())),
});
}
if arguments.is_empty() {
None
} else {
Some(arguments)
}
} else {
None
}
}
codex-rs/vendor/rmcp/src/handler/server/resource.rs
Vendored
+1
View File
@@ -0,0 +1 @@
codex-rs/vendor/rmcp/src/handler/server/router.rs
Vendored
+138
View File
@@ -0,0 +1,138 @@
use std::sync::Arc;
use prompt::{IntoPromptRoute, PromptRoute};
use tool::{IntoToolRoute, ToolRoute};
use super::ServerHandler;
use crate::{
RoleServer, Service,
model::{ClientRequest, ListPromptsResult, ListToolsResult, ServerResult},
service::NotificationContext,
};
pub mod prompt;
pub mod tool;
pub struct Router<S> {
pub tool_router: tool::ToolRouter<S>,
pub prompt_router: prompt::PromptRouter<S>,
pub service: Arc<S>,
}
impl<S> Router<S>
where
S: ServerHandler,
{
pub fn new(service: S) -> Self {
Self {
tool_router: tool::ToolRouter::new(),
prompt_router: prompt::PromptRouter::new(),
service: Arc::new(service),
}
}
pub fn with_tool<R, A>(mut self, route: R) -> Self
where
R: IntoToolRoute<S, A>,
{
self.tool_router.add_route(route.into_tool_route());
self
}
pub fn with_tools(mut self, routes: impl IntoIterator<Item = ToolRoute<S>>) -> Self {
for route in routes {
self.tool_router.add_route(route);
}
self
}
pub fn with_prompt<R, A: 'static>(mut self, route: R) -> Self
where
R: IntoPromptRoute<S, A>,
{
self.prompt_router.add_route(route.into_prompt_route());
self
}
pub fn with_prompts(mut self, routes: impl IntoIterator<Item = PromptRoute<S>>) -> Self {
for route in routes {
self.prompt_router.add_route(route);
}
self
}
}
impl<S> Service<RoleServer> for Router<S>
where
S: ServerHandler,
{
async fn handle_notification(
&self,
notification: <RoleServer as crate::service::ServiceRole>::PeerNot,
context: NotificationContext<RoleServer>,
) -> Result<(), crate::ErrorData> {
self.service
.handle_notification(notification, context)
.await
}
async fn handle_request(
&self,
request: <RoleServer as crate::service::ServiceRole>::PeerReq,
context: crate::service::RequestContext<RoleServer>,
) -> Result<<RoleServer as crate::service::ServiceRole>::Resp, crate::ErrorData> {
match request {
ClientRequest::CallToolRequest(request) => {
if self.tool_router.has_route(request.params.name.as_ref())
|| !self.tool_router.transparent_when_not_found
{
let tool_call_context = crate::handler::server::tool::ToolCallContext::new(
self.service.as_ref(),
request.params,
context,
);
let result = self.tool_router.call(tool_call_context).await?;
Ok(ServerResult::CallToolResult(result))
} else {
self.service
.handle_request(ClientRequest::CallToolRequest(request), context)
.await
}
}
ClientRequest::ListToolsRequest(_) => {
let tools = self.tool_router.list_all();
Ok(ServerResult::ListToolsResult(ListToolsResult {
tools,
..Default::default()
}))
}
ClientRequest::GetPromptRequest(request) => {
if self.prompt_router.has_route(request.params.name.as_ref()) {
let prompt_context = crate::handler::server::prompt::PromptContext::new(
self.service.as_ref(),
request.params.name,
request.params.arguments,
context,
);
let result = self.prompt_router.get_prompt(prompt_context).await?;
Ok(ServerResult::GetPromptResult(result))
} else {
self.service
.handle_request(ClientRequest::GetPromptRequest(request), context)
.await
}
}
ClientRequest::ListPromptsRequest(_) => {
let prompts = self.prompt_router.list_all();
Ok(ServerResult::ListPromptsResult(ListPromptsResult {
prompts,
..Default::default()
}))
}
rest => self.service.handle_request(rest, context).await,
}
}
fn get_info(&self) -> <RoleServer as crate::service::ServiceRole>::Info {
ServerHandler::get_info(&self.service)
}
}
codex-rs/vendor/rmcp/src/handler/server/router/prompt.rs
Vendored
+213
View File
@@ -0,0 +1,213 @@
use std::{borrow::Cow, sync::Arc};
use futures::future::BoxFuture;
use crate::{
handler::server::prompt::{DynGetPromptHandler, GetPromptHandler, PromptContext},
model::{GetPromptResult, Prompt},
};
pub struct PromptRoute<S> {
#[allow(clippy::type_complexity)]
pub get: Arc<DynGetPromptHandler<S>>,
pub attr: crate::model::Prompt,
}
impl<S> std::fmt::Debug for PromptRoute<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PromptRoute")
.field("name", &self.attr.name)
.field("description", &self.attr.description)
.field("arguments", &self.attr.arguments)
.finish()
}
}
impl<S> Clone for PromptRoute<S> {
fn clone(&self) -> Self {
Self {
get: self.get.clone(),
attr: self.attr.clone(),
}
}
}
impl<S: Send + Sync + 'static> PromptRoute<S> {
pub fn new<H, A: 'static>(attr: impl Into<Prompt>, handler: H) -> Self
where
H: GetPromptHandler<S, A> + Send + Sync + Clone + 'static,
{
Self {
get: Arc::new(move |context: PromptContext<S>| {
let handler = handler.clone();
handler.handle(context)
}),
attr: attr.into(),
}
}
pub fn new_dyn<H>(attr: impl Into<Prompt>, handler: H) -> Self
where
H: for<'a> Fn(
PromptContext<'a, S>,
) -> BoxFuture<'a, Result<GetPromptResult, crate::ErrorData>>
+ Send
+ Sync
+ 'static,
{
Self {
get: Arc::new(handler),
attr: attr.into(),
}
}
pub fn name(&self) -> &str {
&self.attr.name
}
}
pub trait IntoPromptRoute<S, A> {
fn into_prompt_route(self) -> PromptRoute<S>;
}
impl<S, H, A, P> IntoPromptRoute<S, A> for (P, H)
where
S: Send + Sync + 'static,
A: 'static,
H: GetPromptHandler<S, A> + Send + Sync + Clone + 'static,
P: Into<Prompt>,
{
fn into_prompt_route(self) -> PromptRoute<S> {
PromptRoute::new(self.0.into(), self.1)
}
}
impl<S> IntoPromptRoute<S, ()> for PromptRoute<S>
where
S: Send + Sync + 'static,
{
fn into_prompt_route(self) -> PromptRoute<S> {
self
}
}
/// Adapter for functions generated by the #\[prompt\] macro
pub struct PromptAttrGenerateFunctionAdapter;
impl<S, F> IntoPromptRoute<S, PromptAttrGenerateFunctionAdapter> for F
where
S: Send + Sync + 'static,
F: Fn() -> PromptRoute<S>,
{
fn into_prompt_route(self) -> PromptRoute<S> {
(self)()
}
}
#[derive(Debug)]
pub struct PromptRouter<S> {
#[allow(clippy::type_complexity)]
pub map: std::collections::HashMap<Cow<'static, str>, PromptRoute<S>>,
}
impl<S> Default for PromptRouter<S> {
fn default() -> Self {
Self {
map: std::collections::HashMap::new(),
}
}
}
impl<S> Clone for PromptRouter<S> {
fn clone(&self) -> Self {
Self {
map: self.map.clone(),
}
}
}
impl<S> IntoIterator for PromptRouter<S> {
type Item = PromptRoute<S>;
type IntoIter = std::collections::hash_map::IntoValues<Cow<'static, str>, PromptRoute<S>>;
fn into_iter(self) -> Self::IntoIter {
self.map.into_values()
}
}
impl<S> PromptRouter<S>
where
S: Send + Sync + 'static,
{
pub fn new() -> Self {
Self {
map: std::collections::HashMap::new(),
}
}
pub fn with_route<R, A: 'static>(mut self, route: R) -> Self
where
R: IntoPromptRoute<S, A>,
{
self.add_route(route.into_prompt_route());
self
}
pub fn add_route(&mut self, item: PromptRoute<S>) {
self.map.insert(item.attr.name.clone().into(), item);
}
pub fn merge(&mut self, other: PromptRouter<S>) {
for item in other.map.into_values() {
self.add_route(item);
}
}
pub fn remove_route(&mut self, name: &str) {
self.map.remove(name);
}
pub fn has_route(&self, name: &str) -> bool {
self.map.contains_key(name)
}
pub async fn get_prompt(
&self,
context: PromptContext<'_, S>,
) -> Result<GetPromptResult, crate::ErrorData> {
let item = self.map.get(context.name.as_str()).ok_or_else(|| {
crate::ErrorData::invalid_params(
format!("prompt '{}' not found", context.name),
Some(serde_json::json!({
"available_prompts": self.list_all().iter().map(|p| &p.name).collect::<Vec<_>>()
})),
)
})?;
(item.get)(context).await
}
pub fn list_all(&self) -> Vec<crate::model::Prompt> {
self.map.values().map(|item| item.attr.clone()).collect()
}
}
impl<S> std::ops::Add<PromptRouter<S>> for PromptRouter<S>
where
S: Send + Sync + 'static,
{
type Output = Self;
fn add(mut self, other: PromptRouter<S>) -> Self::Output {
self.merge(other);
self
}
}
impl<S> std::ops::AddAssign<PromptRouter<S>> for PromptRouter<S>
where
S: Send + Sync + 'static,
{
fn add_assign(&mut self, other: PromptRouter<S>) {
self.merge(other);
}
}
codex-rs/vendor/rmcp/src/handler/server/router/tool.rs
Vendored
+285
View File
@@ -0,0 +1,285 @@
use std::{borrow::Cow, sync::Arc};
use futures::{FutureExt, future::BoxFuture};
use schemars::JsonSchema;
use crate::{
handler::server::{
tool::{CallToolHandler, DynCallToolHandler, ToolCallContext, schema_for_type},
tool_name_validation::validate_and_warn_tool_name,
},
model::{CallToolResult, Tool, ToolAnnotations},
};
pub struct ToolRoute<S> {
#[allow(clippy::type_complexity)]
pub call: Arc<DynCallToolHandler<S>>,
pub attr: crate::model::Tool,
}
impl<S> std::fmt::Debug for ToolRoute<S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ToolRoute")
.field("name", &self.attr.name)
.field("description", &self.attr.description)
.field("input_schema", &self.attr.input_schema)
.finish()
}
}
impl<S> Clone for ToolRoute<S> {
fn clone(&self) -> Self {
Self {
call: self.call.clone(),
attr: self.attr.clone(),
}
}
}
impl<S: Send + Sync + 'static> ToolRoute<S> {
pub fn new<C, A>(attr: impl Into<Tool>, call: C) -> Self
where
C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
{
Self {
call: Arc::new(move |context: ToolCallContext<S>| {
let call = call.clone();
context.invoke(call).boxed()
}),
attr: attr.into(),
}
}
pub fn new_dyn<C>(attr: impl Into<Tool>, call: C) -> Self
where
C: for<'a> Fn(
ToolCallContext<'a, S>,
) -> BoxFuture<'a, Result<CallToolResult, crate::ErrorData>>
+ Send
+ Sync
+ 'static,
{
Self {
call: Arc::new(call),
attr: attr.into(),
}
}
pub fn name(&self) -> &str {
&self.attr.name
}
}
pub trait IntoToolRoute<S, A> {
fn into_tool_route(self) -> ToolRoute<S>;
}
impl<S, C, A, T> IntoToolRoute<S, A> for (T, C)
where
S: Send + Sync + 'static,
C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
T: Into<Tool>,
{
fn into_tool_route(self) -> ToolRoute<S> {
ToolRoute::new(self.0.into(), self.1)
}
}
impl<S> IntoToolRoute<S, ()> for ToolRoute<S>
where
S: Send + Sync + 'static,
{
fn into_tool_route(self) -> ToolRoute<S> {
self
}
}
pub struct ToolAttrGenerateFunctionAdapter;
impl<S, F> IntoToolRoute<S, ToolAttrGenerateFunctionAdapter> for F
where
S: Send + Sync + 'static,
F: Fn() -> ToolRoute<S>,
{
fn into_tool_route(self) -> ToolRoute<S> {
(self)()
}
}
pub trait CallToolHandlerExt<S, A>: Sized
where
Self: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
{
fn name(self, name: impl Into<Cow<'static, str>>) -> WithToolAttr<Self, S, A>;
}
impl<C, S, A> CallToolHandlerExt<S, A> for C
where
C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
{
fn name(self, name: impl Into<Cow<'static, str>>) -> WithToolAttr<Self, S, A> {
WithToolAttr {
attr: Tool::new(
name.into(),
"",
schema_for_type::<crate::model::JsonObject>(),
),
call: self,
_marker: std::marker::PhantomData,
}
}
}
pub struct WithToolAttr<C, S, A>
where
C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
{
pub attr: crate::model::Tool,
pub call: C,
pub _marker: std::marker::PhantomData<fn(S, A)>,
}
impl<C, S, A> IntoToolRoute<S, A> for WithToolAttr<C, S, A>
where
C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
S: Send + Sync + 'static,
{
fn into_tool_route(self) -> ToolRoute<S> {
ToolRoute::new(self.attr, self.call)
}
}
impl<C, S, A> WithToolAttr<C, S, A>
where
C: CallToolHandler<S, A> + Send + Sync + Clone + 'static,
{
pub fn description(mut self, description: impl Into<Cow<'static, str>>) -> Self {
self.attr.description = Some(description.into());
self
}
pub fn parameters<T: JsonSchema + 'static>(mut self) -> Self {
self.attr.input_schema = schema_for_type::<T>();
self
}
pub fn parameters_value(mut self, schema: serde_json::Value) -> Self {
self.attr.input_schema = crate::model::object(schema).into();
self
}
pub fn annotation(mut self, annotation: impl Into<ToolAnnotations>) -> Self {
self.attr.annotations = Some(annotation.into());
self
}
}
#[derive(Debug)]
pub struct ToolRouter<S> {
#[allow(clippy::type_complexity)]
pub map: std::collections::HashMap<Cow<'static, str>, ToolRoute<S>>,
pub transparent_when_not_found: bool,
}
impl<S> Default for ToolRouter<S> {
fn default() -> Self {
Self {
map: std::collections::HashMap::new(),
transparent_when_not_found: false,
}
}
}
impl<S> Clone for ToolRouter<S> {
fn clone(&self) -> Self {
Self {
map: self.map.clone(),
transparent_when_not_found: self.transparent_when_not_found,
}
}
}
impl<S> IntoIterator for ToolRouter<S> {
type Item = ToolRoute<S>;
type IntoIter = std::collections::hash_map::IntoValues<Cow<'static, str>, ToolRoute<S>>;
fn into_iter(self) -> Self::IntoIter {
self.map.into_values()
}
}
impl<S> ToolRouter<S>
where
S: Send + Sync + 'static,
{
pub fn new() -> Self {
Self {
map: std::collections::HashMap::new(),
transparent_when_not_found: false,
}
}
pub fn with_route<R, A>(mut self, route: R) -> Self
where
R: IntoToolRoute<S, A>,
{
self.add_route(route.into_tool_route());
self
}
pub fn add_route(&mut self, item: ToolRoute<S>) {
let new_name = &item.attr.name;
validate_and_warn_tool_name(new_name);
self.map.insert(new_name.clone(), item);
}
pub fn merge(&mut self, other: ToolRouter<S>) {
for item in other.map.into_values() {
self.add_route(item);
}
}
pub fn remove_route(&mut self, name: &str) {
self.map.remove(name);
}
pub fn has_route(&self, name: &str) -> bool {
self.map.contains_key(name)
}
pub async fn call(
&self,
context: ToolCallContext<'_, S>,
) -> Result<CallToolResult, crate::ErrorData> {
let item = self
.map
.get(context.name())
.ok_or_else(|| crate::ErrorData::invalid_params("tool not found", None))?;
let result = (item.call)(context).await?;
Ok(result)
}
pub fn list_all(&self) -> Vec<crate::model::Tool> {
self.map.values().map(|item| item.attr.clone()).collect()
}
/// Get a tool definition by name.
///
/// Returns the tool if found, or `None` if no tool with the given name exists.
pub fn get(&self, name: &str) -> Option<&crate::model::Tool> {
self.map.get(name).map(|r| &r.attr)
}
}
impl<S> std::ops::Add<ToolRouter<S>> for ToolRouter<S>
where
S: Send + Sync + 'static,
{
type Output = Self;
fn add(mut self, other: ToolRouter<S>) -> Self::Output {
self.merge(other);
self
}
}
impl<S> std::ops::AddAssign<ToolRouter<S>> for ToolRouter<S>
where
S: Send + Sync + 'static,
{
fn add_assign(&mut self, other: ToolRouter<S>) {
self.merge(other);
}
}
codex-rs/vendor/rmcp/src/handler/server/tool.rs
Vendored
+333
View File
@@ -0,0 +1,333 @@
use std::{
borrow::Cow,
future::{Future, Ready},
marker::PhantomData,
};
use futures::future::{BoxFuture, FutureExt};
use serde::de::DeserializeOwned;
use super::common::{AsRequestContext, FromContextPart};
pub use super::{
common::{Extension, RequestId, schema_for_output, schema_for_type},
router::tool::{ToolRoute, ToolRouter},
};
use crate::{
RoleServer,
handler::server::wrapper::Parameters,
model::{CallToolRequestParams, CallToolResult, IntoContents, JsonObject},
service::RequestContext,
};
/// Deserialize a JSON object into a type
pub fn parse_json_object<T: DeserializeOwned>(input: JsonObject) -> Result<T, crate::ErrorData> {
serde_json::from_value(serde_json::Value::Object(input)).map_err(|e| {
crate::ErrorData::invalid_params(
format!("failed to deserialize parameters: {error}", error = e),
None,
)
})
}
pub struct ToolCallContext<'s, S> {
pub request_context: RequestContext<RoleServer>,
pub service: &'s S,
pub name: Cow<'static, str>,
pub arguments: Option<JsonObject>,
pub task: Option<JsonObject>,
}
impl<'s, S> ToolCallContext<'s, S> {
pub fn new(
service: &'s S,
CallToolRequestParams {
meta: _,
name,
arguments,
task,
}: CallToolRequestParams,
request_context: RequestContext<RoleServer>,
) -> Self {
Self {
request_context,
service,
name,
arguments,
task,
}
}
pub fn name(&self) -> &str {
&self.name
}
pub fn request_context(&self) -> &RequestContext<RoleServer> {
&self.request_context
}
}
impl<S> AsRequestContext for ToolCallContext<'_, S> {
fn as_request_context(&self) -> &RequestContext<RoleServer> {
&self.request_context
}
fn as_request_context_mut(&mut self) -> &mut RequestContext<RoleServer> {
&mut self.request_context
}
}
pub trait IntoCallToolResult {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData>;
}
impl<T: IntoContents> IntoCallToolResult for T {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
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::ErrorData> {
match self {
Ok(value) => Ok(CallToolResult::success(value.into_contents())),
Err(error) => Ok(CallToolResult::error(error.into_contents())),
}
}
}
impl<T: IntoCallToolResult> IntoCallToolResult for Result<T, crate::ErrorData> {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
match self {
Ok(value) => value.into_call_tool_result(),
Err(error) => Err(error),
}
}
}
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::ErrorData>>,
}
}
}
impl<F, R> Future for IntoCallToolResultFut<F, R>
where
F: Future<Output = R>,
R: IntoCallToolResult,
{
type Output = Result<CallToolResult, crate::ErrorData>;
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::ErrorData> {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
self
}
}
pub trait CallToolHandler<S, A> {
fn call(
self,
context: ToolCallContext<'_, S>,
) -> BoxFuture<'_, Result<CallToolResult, crate::ErrorData>>;
}
pub type DynCallToolHandler<S> = dyn for<'s> Fn(ToolCallContext<'s, S>) -> BoxFuture<'s, Result<CallToolResult, crate::ErrorData>>
+ Send
+ Sync;
// Tool-specific extractor for tool name
pub struct ToolName(pub Cow<'static, str>);
impl<S> FromContextPart<ToolCallContext<'_, S>> for ToolName {
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
Ok(Self(context.name.clone()))
}
}
// Special implementation for Parameters that handles tool arguments
impl<S, P> FromContextPart<ToolCallContext<'_, S>> for Parameters<P>
where
P: DeserializeOwned,
{
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
let arguments = context.arguments.take().unwrap_or_default();
let value: P =
serde_json::from_value(serde_json::Value::Object(arguments)).map_err(|e| {
crate::ErrorData::invalid_params(
format!("failed to deserialize parameters: {error}", error = e),
None,
)
})?;
Ok(Parameters(value))
}
}
// Special implementation for JsonObject that takes tool arguments
impl<S> FromContextPart<ToolCallContext<'_, S>> for JsonObject {
fn from_context_part(context: &mut ToolCallContext<S>) -> Result<Self, crate::ErrorData> {
let object = context.arguments.take().unwrap_or_default();
Ok(object)
}
}
impl<'s, S> ToolCallContext<'s, S> {
pub fn invoke<H, A>(self, h: H) -> BoxFuture<'s, Result<CallToolResult, crate::ErrorData>>
where
H: CallToolHandler<S, A>,
{
h.call(self)
}
}
#[allow(clippy::type_complexity)]
pub struct AsyncAdapter<P, Fut, R>(PhantomData<fn(P) -> fn(Fut) -> R>);
pub struct SyncAdapter<P, R>(PhantomData<fn(P) -> R>);
// #[allow(clippy::type_complexity)]
pub struct AsyncMethodAdapter<P, R>(PhantomData<fn(P) -> R>);
pub struct SyncMethodAdapter<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<$($Tn,)* S, F, R> CallToolHandler<S, AsyncMethodAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> ,
)*
F: FnOnce(&S, $($Tn,)*) -> BoxFuture<'_, R>,
// Need RTN support here(I guess), https://github.com/rust-lang/rust/pull/138424
// Fut: Future<Output = R> + Send + 'a,
R: IntoCallToolResult + Send + 'static,
S: Send + Sync + 'static,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn call(
self,
mut context: ToolCallContext<'_, S>,
) -> BoxFuture<'_, Result<CallToolResult, crate::ErrorData>>{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let service = context.service;
let fut = self(service, $($Tn,)*);
async move {
let result = fut.await;
result.into_call_tool_result()
}.boxed()
}
}
impl<$($Tn,)* S, F, Fut, R> CallToolHandler<S, AsyncAdapter<($($Tn,)*), Fut, R>> for F
where
$(
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> ,
)*
F: FnOnce($($Tn,)*) -> Fut + Send + ,
Fut: Future<Output = R> + Send + 'static,
R: IntoCallToolResult + Send + 'static,
S: Send + Sync,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn call(
self,
mut context: ToolCallContext<S>,
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>>{
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
let fut = self($($Tn,)*);
async move {
let result = fut.await;
result.into_call_tool_result()
}.boxed()
}
}
impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncMethodAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> + ,
)*
F: FnOnce(&S, $($Tn,)*) -> R + Send + ,
R: IntoCallToolResult + Send + ,
S: Send + Sync,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn call(
self,
mut context: ToolCallContext<S>,
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
std::future::ready(self(context.service, $($Tn,)*).into_call_tool_result()).boxed()
}
}
impl<$($Tn,)* S, F, R> CallToolHandler<S, SyncAdapter<($($Tn,)*), R>> for F
where
$(
$Tn: for<'a> FromContextPart<ToolCallContext<'a, S>> + ,
)*
F: FnOnce($($Tn,)*) -> R + Send + ,
R: IntoCallToolResult + Send + ,
S: Send + Sync,
{
#[allow(unused_variables, non_snake_case, unused_mut)]
fn call(
self,
mut context: ToolCallContext<S>,
) -> BoxFuture<'static, Result<CallToolResult, crate::ErrorData>> {
$(
let result = $Tn::from_context_part(&mut context);
let $Tn = match result {
Ok(value) => value,
Err(e) => return std::future::ready(Err(e)).boxed(),
};
)*
std::future::ready(self($($Tn,)*).into_call_tool_result()).boxed()
}
}
};
}
impl_for!(T0 T1 T2 T3 T4 T5 T6 T7 T8 T9 T10 T11 T12 T13 T14 T15);
codex-rs/vendor/rmcp/src/handler/server/tool_name_validation.rs
Vendored
+261
View File
@@ -0,0 +1,261 @@
//! Tool name validation utilities according to SEP: Specify Format for Tool Names
//!
//! Tool names SHOULD be between 1 and 128 characters in length (inclusive).
//! Tool names are case-sensitive.
//! Allowed characters: uppercase and lowercase ASCII letters (A-Z, a-z), digits
//! (0-9), underscore (_), dash (-), and dot (.).
//! Tool names SHOULD NOT contain spaces, commas, or other special characters.
use std::collections::HashSet;
/// Result of tool name validation containing validation status and warnings.
#[derive(Debug, Clone, PartialEq, Eq)]
struct ToolNameValidationResult {
/// Whether the tool name is valid according to the specification
is_valid: bool,
/// Array of warning messages about non-conforming aspects of the tool name
warnings: Vec<String>,
}
impl ToolNameValidationResult {
/// Create a new validation result
fn new(is_valid: bool, warnings: Vec<String>) -> Self {
Self { is_valid, warnings }
}
}
/// Validates a tool name according to the SEP specification.
fn validate_tool_name(name: &str) -> ToolNameValidationResult {
let mut warnings = Vec::new();
// Check length
if name.is_empty() {
return ToolNameValidationResult::new(false, vec!["Tool name cannot be empty".to_string()]);
}
if name.len() > 128 {
return ToolNameValidationResult::new(
false,
vec![format!(
"Tool name exceeds maximum length of 128 characters (current: {})",
name.len()
)],
);
}
// Check for specific problematic patterns (these are warnings, not validation failures)
if name.contains(' ') {
warnings.push("Tool name contains spaces, which may cause parsing issues".to_string());
}
if name.contains(',') {
warnings.push("Tool name contains commas, which may cause parsing issues".to_string());
}
// Check for potentially confusing patterns (leading/trailing dashes, dots, slashes)
if name.starts_with('-') || name.ends_with('-') {
warnings.push(
"Tool name starts or ends with a dash, which may cause parsing issues in some contexts"
.to_string(),
);
}
if name.starts_with('.') || name.ends_with('.') {
warnings.push(
"Tool name starts or ends with a dot, which may cause parsing issues in some contexts"
.to_string(),
);
}
// Check for invalid characters
let mut invalid_chars = HashSet::new();
let valid_chars: HashSet<char> =
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._-"
.chars()
.collect();
for ch in name.chars() {
if !valid_chars.contains(&ch) {
invalid_chars.insert(ch);
}
}
if !invalid_chars.is_empty() {
let invalid_chars_list: Vec<String> =
invalid_chars.iter().map(|c| format!("\"{}\"", c)).collect();
warnings.push(format!(
"Tool name contains invalid characters: {}",
invalid_chars_list.join(", ")
));
warnings.push(
"Allowed characters are: A-Z, a-z, 0-9, underscore (_), dash (-), and dot (.)"
.to_string(),
);
return ToolNameValidationResult::new(false, warnings);
}
// Verify the pattern matches (double check with character-by-character validation)
// We've already validated characters above, just need to verify length is within bounds
if name.is_empty() || name.len() > 128 {
return ToolNameValidationResult::new(
false,
vec!["Tool name length must be between 1 and 128 characters".to_string()],
);
}
ToolNameValidationResult::new(true, warnings)
}
/// Issues warnings for non-conforming tool names.
fn issue_tool_name_warning(name: &str, warnings: &[String]) {
tracing::warn!("Tool name validation warning for \"{}\":", name);
for warning in warnings {
tracing::warn!(" - {}", warning);
}
tracing::warn!("Tool registration will proceed, but this may cause compatibility issues.");
tracing::warn!("Consider updating the tool name to conform to the MCP tool naming standard.");
tracing::warn!(
"See SEP: Specify Format for Tool Names (https://github.com/modelcontextprotocol/modelcontextprotocol/issues/986) for more details."
);
}
/// Validates a tool name and issues warnings for non-conforming names.
pub fn validate_and_warn_tool_name(name: &str) -> bool {
let result = validate_tool_name(name);
if !result.warnings.is_empty() {
issue_tool_name_warning(name, &result.warnings);
}
result.is_valid
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_valid_tool_names() {
let max_length_name = "a".repeat(128);
let valid_names = vec![
"my_tool",
"MyTool",
"my-tool",
"my.tool",
"tool123",
"a",
max_length_name.as_str(), // Maximum length
];
for name in valid_names {
let result = validate_tool_name(name);
assert!(result.is_valid, "Tool name '{}' should be valid", name);
}
}
#[test]
fn test_empty_tool_name() {
let result = validate_tool_name("");
assert!(!result.is_valid);
assert!(
result
.warnings
.contains(&"Tool name cannot be empty".to_string())
);
}
#[test]
fn test_too_long_tool_name() {
let name = "a".repeat(129);
let result = validate_tool_name(&name);
assert!(!result.is_valid);
assert!(result.warnings[0].contains("exceeds maximum length"));
}
#[test]
fn test_tool_name_with_spaces() {
let result = validate_tool_name("my tool");
assert!(!result.is_valid);
assert!(
result
.warnings
.iter()
.any(|w| w.contains("contains spaces"))
);
}
#[test]
fn test_tool_name_with_commas() {
let result = validate_tool_name("my,tool");
assert!(!result.is_valid);
assert!(
result
.warnings
.iter()
.any(|w| w.contains("contains commas"))
);
}
#[test]
fn test_tool_name_starting_with_dash() {
let result = validate_tool_name("-tool");
assert!(result.is_valid); // Still valid, but has warning
assert!(
result
.warnings
.iter()
.any(|w| w.contains("starts or ends with a dash"))
);
}
#[test]
fn test_tool_name_ending_with_dot() {
let result = validate_tool_name("tool.");
assert!(result.is_valid); // Still valid, but has warning
assert!(
result
.warnings
.iter()
.any(|w| w.contains("starts or ends with a dot"))
);
}
#[test]
fn test_tool_name_with_invalid_characters() {
let result = validate_tool_name("my@tool");
assert!(!result.is_valid);
assert!(
result
.warnings
.iter()
.any(|w| w.contains("contains invalid characters"))
);
}
#[test]
fn test_tool_name_all_special_characters_allowed() {
let valid_chars = vec!['_', '-', '.'];
for ch in valid_chars {
let name = format!("tool{}", ch);
let result = validate_tool_name(&name);
assert!(
result.is_valid,
"Tool name with character '{}' should be valid",
ch
);
}
}
#[test]
fn test_minimum_length() {
let result = validate_tool_name("a");
assert!(result.is_valid);
}
#[test]
fn test_maximum_length() {
let name = "a".repeat(128);
let result = validate_tool_name(&name);
assert!(result.is_valid);
}
}
codex-rs/vendor/rmcp/src/handler/server/wrapper.rs
Vendored
+4
View File
@@ -0,0 +1,4 @@
mod json;
mod parameters;
pub use json::*;
pub use parameters::*;
codex-rs/vendor/rmcp/src/handler/server/wrapper/json.rs
Vendored
+54
View File
@@ -0,0 +1,54 @@
use std::borrow::Cow;
use schemars::JsonSchema;
use serde::Serialize;
use crate::{
handler::server::tool::IntoCallToolResult,
model::{CallToolResult, IntoContents},
};
/// Json wrapper for structured output
///
/// When used with tools, this wrapper indicates that the value should be
/// serialized as structured JSON content with an associated schema.
/// The framework will place the JSON in the `structured_content` field
/// of the tool result rather than the regular `content` field.
pub struct Json<T>(pub T);
// Implement JsonSchema for Json<T> to delegate to T's schema
impl<T: JsonSchema> JsonSchema for Json<T> {
fn schema_name() -> Cow<'static, str> {
T::schema_name()
}
fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
T::json_schema(generator)
}
}
// Implementation for Json<T> to create structured content
impl<T: Serialize + JsonSchema + 'static> IntoCallToolResult for Json<T> {
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
let value = serde_json::to_value(self.0).map_err(|e| {
crate::ErrorData::internal_error(
format!("Failed to serialize structured content: {}", e),
None,
)
})?;
Ok(CallToolResult::structured(value))
}
}
// Implementation for Result<Json<T>, E>
impl<T: Serialize + JsonSchema + 'static, E: IntoContents> IntoCallToolResult
for Result<Json<T>, E>
{
fn into_call_tool_result(self) -> Result<CallToolResult, crate::ErrorData> {
match self {
Ok(value) => value.into_call_tool_result(),
Err(error) => Ok(CallToolResult::error(error.into_contents())),
}
}
}
codex-rs/vendor/rmcp/src/handler/server/wrapper/parameters.rs
Vendored
+55
View File
@@ -0,0 +1,55 @@
use schemars::JsonSchema;
/// Parameter extractor for tools and prompts
///
/// When used in tool and prompt handlers, this wrapper extracts and deserializes
/// parameters from the incoming request. The framework will automatically parse
/// the JSON arguments from tool calls or prompt arguments and deserialize them
/// into the specified type `P`.
///
/// The `#[serde(transparent)]` attribute ensures that the wrapper doesn't add
/// an extra layer in the JSON structure - it directly delegates serialization
/// and deserialization to the inner type `P`.
///
/// # Usage
///
/// Use `Parameters<T>` as a parameter in your tool or prompt handler functions:
///
/// ```rust
/// # use rmcp::handler::server::wrapper::Parameters;
/// # use schemars::JsonSchema;
/// # use serde::{Deserialize, Serialize};
/// #[derive(Deserialize, JsonSchema)]
/// struct CalculationRequest {
/// operation: String,
/// a: f64,
/// b: f64,
/// }
///
/// // In a tool handler
/// async fn calculate(params: Parameters<CalculationRequest>) -> Result<String, String> {
/// let request = params.0; // Extract the inner value
/// match request.operation.as_str() {
/// "add" => Ok((request.a + request.b).to_string()),
/// _ => Err("Unknown operation".to_string()),
/// }
/// }
/// ```
///
/// The framework handles the extraction automatically:
/// - For tools: Parses the `arguments` field from tool call requests
/// - For prompts: Parses the `arguments` field from prompt requests
/// - Returns appropriate error responses if deserialization fails
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(transparent)]
pub struct Parameters<P>(pub P);
impl<P: JsonSchema> JsonSchema for Parameters<P> {
fn schema_name() -> std::borrow::Cow<'static, str> {
P::schema_name()
}
fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
P::json_schema(generator)
}
}
codex-rs/vendor/rmcp/src/lib.rs
Vendored
+52
View File
@@ -0,0 +1,52 @@
#![cfg_attr(docsrs, feature(doc_cfg))]
#![cfg_attr(docsrs, allow(unused_attributes))]
#![doc = include_str!("../README.md")]
mod error;
#[allow(deprecated)]
pub use error::{Error, ErrorData, RmcpError};
/// Basic data types in MCP specification
pub mod model;
#[cfg(any(feature = "client", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(any(feature = "client", feature = "server"))))]
pub mod service;
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub use handler::client::ClientHandler;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub use handler::server::ServerHandler;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub use handler::server::wrapper::Json;
#[cfg(any(feature = "client", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(any(feature = "client", feature = "server"))))]
pub use service::{Peer, Service, ServiceError, ServiceExt};
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub use service::{RoleClient, serve_client};
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub use service::{RoleServer, serve_server};
pub mod handler;
pub mod task_manager;
pub mod transport;
// re-export
#[cfg(all(feature = "macros", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(all(feature = "macros", feature = "server"))))]
pub use pastey::paste;
#[cfg(all(feature = "macros", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(all(feature = "macros", feature = "server"))))]
pub use rmcp_macros::*;
#[cfg(any(feature = "macros", feature = "server"))]
#[cfg_attr(docsrs, doc(cfg(any(feature = "macros", feature = "server"))))]
pub use schemars;
#[cfg(feature = "macros")]
#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
pub use serde;
#[cfg(feature = "macros")]
#[cfg_attr(docsrs, doc(cfg(feature = "macros")))]
pub use serde_json;
codex-rs/vendor/rmcp/src/model.rs
Vendored
+3352
View File
File diff suppressed because it is too large Load Diff
codex-rs/vendor/rmcp/src/model/annotated.rs
Vendored
+224
View File
@@ -0,0 +1,224 @@
use std::ops::{Deref, DerefMut};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use super::{
RawAudioContent, RawContent, RawEmbeddedResource, RawImageContent, RawResource,
RawResourceTemplate, RawTextContent, Role,
};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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", rename = "lastModified")]
pub last_modified: 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),
last_modified: Some(timestamp),
audience: None,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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.last_modified)
}
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,
last_modified: 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),
last_modified: 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 {
last_modified: Some(timestamp),
..annotations
}),
}
} else {
Annotated {
raw: self.raw,
annotations: Some(Annotations {
last_modified: 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!(RawAudioContent);
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 {
last_modified: Some(timestamp),
..Default::default()
})
}
fn with_timestamp_now(self) -> Annotated<Self>
where
Self: Sized,
{
self.with_timestamp(Utc::now())
}
}
codex-rs/vendor/rmcp/src/model/capabilities.rs
Vendored
+731
View File
@@ -0,0 +1,731 @@
use std::{collections::BTreeMap, marker::PhantomData};
use pastey::paste;
use serde::{Deserialize, Serialize};
use super::JsonObject;
pub type ExperimentalCapabilities = BTreeMap<String, JsonObject>;
/// MCP extension capabilities map.
///
/// Keys are extension identifiers in the format `{vendor-prefix}/{extension-name}`
/// (e.g., `io.modelcontextprotocol/ui`, `io.modelcontextprotocol/oauth-client-credentials`).
/// Values are per-extension settings objects. An empty object indicates support with no settings.
///
/// # Example
///
/// ```rust
/// use rmcp::model::ExtensionCapabilities;
/// use serde_json::json;
///
/// let mut extensions = ExtensionCapabilities::new();
/// extensions.insert(
/// "io.modelcontextprotocol/ui".to_string(),
/// serde_json::from_value(json!({
/// "mimeTypes": ["text/html;profile=mcp-app"]
/// })).unwrap()
/// );
/// ```
pub type ExtensionCapabilities = BTreeMap<String, JsonObject>;
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RootsCapabilities {
#[serde(skip_serializing_if = "Option::is_none")]
pub list_changed: Option<bool>,
}
/// Task capabilities shared by client and server.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct TasksCapability {
#[serde(skip_serializing_if = "Option::is_none")]
pub requests: Option<TaskRequestsCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub list: Option<JsonObject>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cancel: Option<JsonObject>,
}
/// Request types that support task-augmented execution.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct TaskRequestsCapability {
#[serde(skip_serializing_if = "Option::is_none")]
pub sampling: Option<SamplingTaskCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub elicitation: Option<ElicitationTaskCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<ToolsTaskCapability>,
}
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct SamplingTaskCapability {
#[serde(skip_serializing_if = "Option::is_none")]
pub create_message: Option<JsonObject>,
}
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ElicitationTaskCapability {
#[serde(skip_serializing_if = "Option::is_none")]
pub create: Option<JsonObject>,
}
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ToolsTaskCapability {
#[serde(skip_serializing_if = "Option::is_none")]
pub call: Option<JsonObject>,
}
impl TasksCapability {
/// Default client tasks capability with sampling and elicitation support.
pub fn client_default() -> Self {
Self {
list: Some(JsonObject::new()),
cancel: Some(JsonObject::new()),
requests: Some(TaskRequestsCapability {
sampling: Some(SamplingTaskCapability {
create_message: Some(JsonObject::new()),
}),
elicitation: Some(ElicitationTaskCapability {
create: Some(JsonObject::new()),
}),
tools: None,
}),
}
}
/// Default server tasks capability with tools/call support.
pub fn server_default() -> Self {
Self {
list: Some(JsonObject::new()),
cancel: Some(JsonObject::new()),
requests: Some(TaskRequestsCapability {
sampling: None,
elicitation: None,
tools: Some(ToolsTaskCapability {
call: Some(JsonObject::new()),
}),
}),
}
}
pub fn supports_list(&self) -> bool {
self.list.is_some()
}
pub fn supports_cancel(&self) -> bool {
self.cancel.is_some()
}
pub fn supports_tools_call(&self) -> bool {
self.requests
.as_ref()
.and_then(|r| r.tools.as_ref())
.and_then(|t| t.call.as_ref())
.is_some()
}
pub fn supports_sampling_create_message(&self) -> bool {
self.requests
.as_ref()
.and_then(|r| r.sampling.as_ref())
.and_then(|s| s.create_message.as_ref())
.is_some()
}
pub fn supports_elicitation_create(&self) -> bool {
self.requests
.as_ref()
.and_then(|r| r.elicitation.as_ref())
.and_then(|e| e.create.as_ref())
.is_some()
}
}
/// Capability for handling elicitation requests from servers.
/// Elicitation allows servers to request interactive input from users during tool execution.
/// This capability indicates that a client can handle elicitation requests and present
/// appropriate UI to users for collecting the requested information.
///
/// Capability for form mode elicitation.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct FormElicitationCapability {
/// Whether the client supports JSON Schema validation for elicitation responses.
/// When true, the client will validate user input against the requested_schema
/// before sending the response back to the server.
#[serde(skip_serializing_if = "Option::is_none")]
pub schema_validation: Option<bool>,
}
/// Capability for URL mode elicitation.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct UrlElicitationCapability {}
/// Elicitation allows servers to request interactive input from users during tool execution.
/// This capability indicates that a client can handle elicitation requests and present
/// appropriate UI to users for collecting the requested information.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ElicitationCapability {
/// Whether client supports form-based elicitation.
#[serde(skip_serializing_if = "Option::is_none")]
pub form: Option<FormElicitationCapability>,
/// Whether client supports URL-based elicitation.
#[serde(skip_serializing_if = "Option::is_none")]
pub url: Option<UrlElicitationCapability>,
}
/// Sampling capability with optional sub-capabilities (SEP-1577).
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct SamplingCapability {
/// Support for `tools` and `toolChoice` parameters
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<JsonObject>,
/// Support for `includeContext` (soft-deprecated)
#[serde(skip_serializing_if = "Option::is_none")]
pub context: Option<JsonObject>,
}
///
/// # 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)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ClientCapabilities {
#[serde(skip_serializing_if = "Option::is_none")]
pub experimental: Option<ExperimentalCapabilities>,
/// Optional MCP extensions that the client supports (SEP-1724).
/// Keys are extension identifiers (e.g., `"io.modelcontextprotocol/ui"`),
/// values are per-extension settings objects. An empty object indicates
/// support with no settings.
#[serde(skip_serializing_if = "Option::is_none")]
pub extensions: Option<ExtensionCapabilities>,
#[serde(skip_serializing_if = "Option::is_none")]
pub roots: Option<RootsCapabilities>,
/// Capability for LLM sampling requests (SEP-1577)
#[serde(skip_serializing_if = "Option::is_none")]
pub sampling: Option<SamplingCapability>,
/// Capability to handle elicitation requests from servers for interactive user input
#[serde(skip_serializing_if = "Option::is_none")]
pub elicitation: Option<ElicitationCapability>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tasks: Option<TasksCapability>,
}
///
/// ## 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")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ServerCapabilities {
#[serde(skip_serializing_if = "Option::is_none")]
pub experimental: Option<ExperimentalCapabilities>,
/// Optional MCP extensions that the server supports (SEP-1724).
/// Keys are extension identifiers (e.g., `"io.modelcontextprotocol/apps"`),
/// values are per-extension settings objects. An empty object indicates
/// support with no settings.
#[serde(skip_serializing_if = "Option::is_none")]
pub extensions: Option<ExtensionCapabilities>,
#[serde(skip_serializing_if = "Option::is_none")]
pub logging: Option<JsonObject>,
#[serde(skip_serializing_if = "Option::is_none")]
pub completions: 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>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tasks: Option<TasksCapability>,
}
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,
extensions: ExtensionCapabilities,
logging: JsonObject,
completions: JsonObject,
prompts: PromptsCapability,
resources: ResourcesCapability,
tools: ToolsCapability,
tasks: TasksCapability
}
}
impl<
const E: bool,
const EXT: bool,
const L: bool,
const C: bool,
const P: bool,
const R: bool,
const TASKS: bool,
> ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, EXT, L, C, P, R, true, TASKS>>
{
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 EXT: bool,
const L: bool,
const C: bool,
const R: bool,
const T: bool,
const TASKS: bool,
> ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, EXT, L, C, true, R, T, TASKS>>
{
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 EXT: bool,
const L: bool,
const C: bool,
const P: bool,
const T: bool,
const TASKS: bool,
> ServerCapabilitiesBuilder<ServerCapabilitiesBuilderState<E, EXT, L, C, P, true, T, TASKS>>
{
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,
extensions: ExtensionCapabilities,
roots: RootsCapabilities,
sampling: SamplingCapability,
elicitation: ElicitationCapability,
tasks: TasksCapability,
}
}
impl<const E: bool, const EXT: bool, const S: bool, const EL: bool, const TASKS: bool>
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, EXT, true, S, EL, TASKS>>
{
pub fn enable_roots_list_changed(mut self) -> Self {
if let Some(c) = self.roots.as_mut() {
c.list_changed = Some(true);
}
self
}
}
impl<const E: bool, const EXT: bool, const R: bool, const EL: bool, const TASKS: bool>
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, EXT, R, true, EL, TASKS>>
{
/// Enable tool calling in sampling requests
pub fn enable_sampling_tools(mut self) -> Self {
if let Some(c) = self.sampling.as_mut() {
c.tools = Some(JsonObject::default());
}
self
}
/// Enable context inclusion in sampling (soft-deprecated)
pub fn enable_sampling_context(mut self) -> Self {
if let Some(c) = self.sampling.as_mut() {
c.context = Some(JsonObject::default());
}
self
}
}
#[cfg(feature = "elicitation")]
impl<const E: bool, const EXT: bool, const R: bool, const S: bool, const TASKS: bool>
ClientCapabilitiesBuilder<ClientCapabilitiesBuilderState<E, EXT, R, S, true, TASKS>>
{
/// Enable JSON Schema validation for elicitation responses in form mode.
/// When enabled, the client will validate user input against the requested_schema
/// before sending responses back to the server.
pub fn enable_elicitation_schema_validation(mut self) -> Self {
if let Some(c) = self.elicitation.as_mut() {
c.form = Some(FormElicitationCapability {
schema_validation: 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),
})
);
}
#[test]
fn test_task_capabilities_deserialization() {
// Test deserializing from the MCP spec format
let json = serde_json::json!({
"list": {},
"cancel": {},
"requests": {
"tools": { "call": {} }
}
});
let tasks: TasksCapability = serde_json::from_value(json).unwrap();
assert!(tasks.list.is_some());
assert!(tasks.cancel.is_some());
assert!(tasks.requests.is_some());
let requests = tasks.requests.unwrap();
assert!(requests.tools.is_some());
assert!(requests.tools.unwrap().call.is_some());
}
#[test]
fn test_tasks_capability_client_default() {
let tasks = TasksCapability::client_default();
// Verify structure
assert!(tasks.supports_list());
assert!(tasks.supports_cancel());
assert!(tasks.supports_sampling_create_message());
assert!(tasks.supports_elicitation_create());
assert!(!tasks.supports_tools_call());
// Verify serialization matches expected format
let json = serde_json::to_value(&tasks).unwrap();
assert_eq!(json["list"], serde_json::json!({}));
assert_eq!(json["cancel"], serde_json::json!({}));
assert_eq!(
json["requests"]["sampling"]["createMessage"],
serde_json::json!({})
);
assert_eq!(
json["requests"]["elicitation"]["create"],
serde_json::json!({})
);
}
#[test]
fn test_tasks_capability_server_default() {
let tasks = TasksCapability::server_default();
// Verify structure
assert!(tasks.supports_list());
assert!(tasks.supports_cancel());
assert!(tasks.supports_tools_call());
assert!(!tasks.supports_sampling_create_message());
assert!(!tasks.supports_elicitation_create());
// Verify serialization matches expected format
let json = serde_json::to_value(&tasks).unwrap();
assert_eq!(json["list"], serde_json::json!({}));
assert_eq!(json["cancel"], serde_json::json!({}));
assert_eq!(json["requests"]["tools"]["call"], serde_json::json!({}));
}
#[test]
fn test_client_extensions_capability() {
// Test building ClientCapabilities with extensions (MCP Apps support)
let mut extensions = ExtensionCapabilities::new();
extensions.insert(
"io.modelcontextprotocol/ui".to_string(),
serde_json::from_value(serde_json::json!({
"mimeTypes": ["text/html;profile=mcp-app"]
}))
.unwrap(),
);
let capabilities = ClientCapabilities::builder()
.enable_extensions_with(extensions)
.enable_sampling()
.build();
// Verify serialization matches MCP Apps spec format
let json = serde_json::to_value(&capabilities).unwrap();
assert_eq!(
json["extensions"]["io.modelcontextprotocol/ui"]["mimeTypes"],
serde_json::json!(["text/html;profile=mcp-app"])
);
assert!(json["sampling"].is_object());
}
#[test]
fn test_server_extensions_capability() {
// Test building ServerCapabilities with extensions
let mut extensions = ExtensionCapabilities::new();
extensions.insert(
"io.modelcontextprotocol/apps".to_string(),
serde_json::from_value(serde_json::json!({})).unwrap(),
);
let capabilities = ServerCapabilities::builder()
.enable_extensions_with(extensions)
.enable_tools()
.build();
// Verify serialization
let json = serde_json::to_value(&capabilities).unwrap();
assert!(json["extensions"]["io.modelcontextprotocol/apps"].is_object());
assert!(json["tools"].is_object());
}
#[test]
fn test_extensions_deserialization() {
// Test deserializing capabilities with extensions from JSON
let json = serde_json::json!({
"extensions": {
"io.modelcontextprotocol/ui": {
"mimeTypes": ["text/html;profile=mcp-app"]
}
},
"sampling": {}
});
let capabilities: ClientCapabilities = serde_json::from_value(json).unwrap();
assert!(capabilities.extensions.is_some());
let extensions = capabilities.extensions.unwrap();
assert!(extensions.contains_key("io.modelcontextprotocol/ui"));
let ui_ext = extensions.get("io.modelcontextprotocol/ui").unwrap();
assert!(ui_ext.contains_key("mimeTypes"));
}
#[test]
fn test_extensions_empty_settings() {
// Test that empty extension settings work (indicates support with no settings)
let mut extensions = ExtensionCapabilities::new();
extensions.insert(
"io.modelcontextprotocol/oauth-client-credentials".to_string(),
JsonObject::new(),
);
let capabilities = ClientCapabilities::builder()
.enable_extensions_with(extensions)
.build();
let json = serde_json::to_value(&capabilities).unwrap();
assert_eq!(
json["extensions"]["io.modelcontextprotocol/oauth-client-credentials"],
serde_json::json!({})
);
}
}
codex-rs/vendor/rmcp/src/model/content.rs
Vendored
+364
View File
@@ -0,0 +1,364 @@
//! 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 serde::{Deserialize, Serialize};
use serde_json::json;
use super::{AnnotateAble, Annotated, resource::ResourceContents};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RawTextContent {
pub text: String,
/// Optional protocol-level metadata for this content block
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<super::Meta>,
}
pub type TextContent = Annotated<RawTextContent>;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RawImageContent {
/// The base64-encoded image
pub data: String,
pub mime_type: String,
/// Optional protocol-level metadata for this content block
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<super::Meta>,
}
pub type ImageContent = Annotated<RawImageContent>;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RawEmbeddedResource {
/// Optional protocol-level metadata for this content block
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<super::Meta>,
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(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RawAudioContent {
pub data: String,
pub mime_type: String,
}
pub type AudioContent = Annotated<RawAudioContent>;
/// Tool call request from assistant (SEP-1577).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ToolUseContent {
/// Unique identifier for this tool call
pub id: String,
/// Name of the tool to call
pub name: String,
/// Input arguments for the tool
pub input: super::JsonObject,
/// Optional metadata (preserved for caching)
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<super::Meta>,
}
/// Tool execution result in user message (SEP-1577).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ToolResultContent {
/// Optional metadata
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<super::Meta>,
/// ID of the corresponding tool use
pub tool_use_id: String,
/// Content blocks returned by the tool
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub content: Vec<Content>,
/// Optional structured result
#[serde(skip_serializing_if = "Option::is_none")]
pub structured_content: Option<super::JsonObject>,
/// Whether tool execution failed
#[serde(skip_serializing_if = "Option::is_none")]
pub is_error: Option<bool>,
}
impl ToolUseContent {
pub fn new(id: impl Into<String>, name: impl Into<String>, input: super::JsonObject) -> Self {
Self {
id: id.into(),
name: name.into(),
input,
meta: None,
}
}
}
impl ToolResultContent {
pub fn new(tool_use_id: impl Into<String>, content: Vec<Content>) -> Self {
Self {
meta: None,
tool_use_id: tool_use_id.into(),
content,
structured_content: None,
is_error: None,
}
}
pub fn error(tool_use_id: impl Into<String>, content: Vec<Content>) -> Self {
Self {
meta: None,
tool_use_id: tool_use_id.into(),
content,
structured_content: None,
is_error: Some(true),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub enum RawContent {
Text(RawTextContent),
Image(RawImageContent),
Resource(RawEmbeddedResource),
Audio(RawAudioContent),
ResourceLink(super::resource::RawResource),
}
pub type Content = Annotated<RawContent>;
impl RawContent {
pub fn json<S: Serialize>(json: S) -> Result<Self, crate::ErrorData> {
let json = serde_json::to_string(&json).map_err(|e| {
crate::ErrorData::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(),
meta: None,
})
}
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(),
meta: None,
})
}
pub fn resource(resource: ResourceContents) -> Self {
RawContent::Resource(RawEmbeddedResource {
meta: None,
resource,
})
}
pub fn embedded_text<S: Into<String>, T: Into<String>>(uri: S, content: T) -> Self {
RawContent::Resource(RawEmbeddedResource {
meta: None,
resource: ResourceContents::TextResourceContents {
uri: uri.into(),
mime_type: Some("text".to_string()),
text: content.into(),
meta: None,
},
})
}
/// 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,
}
}
/// Get the resource link if this is a ResourceLink variant
pub fn as_resource_link(&self) -> Option<&super::resource::RawResource> {
match self {
RawContent::ResourceLink(link) => Some(link),
_ => None,
}
}
/// Create a resource link content
pub fn resource_link(resource: super::resource::RawResource) -> Self {
RawContent::ResourceLink(resource)
}
}
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::ErrorData> {
RawContent::json(json).map(|c| c.no_annotation())
}
/// Create a resource link content
pub fn resource_link(resource: super::resource::RawResource) -> Self {
RawContent::resource_link(resource).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)]
}
}
impl IntoContents for () {
fn into_contents(self) -> Vec<Content> {
vec![]
}
}
#[cfg(test)]
mod tests {
use serde_json;
use super::*;
#[test]
fn test_image_content_serialization() {
let image_content = RawImageContent {
data: "base64data".to_string(),
mime_type: "image/png".to_string(),
meta: None,
};
let json = serde_json::to_string(&image_content).unwrap();
println!("ImageContent JSON: {}", json);
// Verify it contains mimeType (camelCase) not mime_type (snake_case)
assert!(json.contains("mimeType"));
assert!(!json.contains("mime_type"));
}
#[test]
fn test_audio_content_serialization() {
let audio_content = RawAudioContent {
data: "base64audiodata".to_string(),
mime_type: "audio/wav".to_string(),
};
let json = serde_json::to_string(&audio_content).unwrap();
println!("AudioContent JSON: {}", json);
// Verify it contains mimeType (camelCase) not mime_type (snake_case)
assert!(json.contains("mimeType"));
assert!(!json.contains("mime_type"));
}
#[test]
fn test_resource_link_serialization() {
use super::super::resource::RawResource;
let resource_link = RawContent::ResourceLink(RawResource {
uri: "file:///test.txt".to_string(),
name: "test.txt".to_string(),
title: None,
description: Some("A test file".to_string()),
mime_type: Some("text/plain".to_string()),
size: Some(100),
icons: None,
meta: None,
});
let json = serde_json::to_string(&resource_link).unwrap();
println!("ResourceLink JSON: {}", json);
// Verify it contains the correct type tag
assert!(json.contains("\"type\":\"resource_link\""));
assert!(json.contains("\"uri\":\"file:///test.txt\""));
assert!(json.contains("\"name\":\"test.txt\""));
}
#[test]
fn test_resource_link_deserialization() {
let json = r#"{
"type": "resource_link",
"uri": "file:///example.txt",
"name": "example.txt",
"description": "Example file",
"mimeType": "text/plain"
}"#;
let content: RawContent = serde_json::from_str(json).unwrap();
if let RawContent::ResourceLink(resource) = content {
assert_eq!(resource.uri, "file:///example.txt");
assert_eq!(resource.name, "example.txt");
assert_eq!(resource.description, Some("Example file".to_string()));
assert_eq!(resource.mime_type, Some("text/plain".to_string()));
} else {
panic!("Expected ResourceLink variant");
}
}
}
codex-rs/vendor/rmcp/src/model/elicitation_schema.rs
Vendored
+2010
View File
File diff suppressed because it is too large Load Diff
codex-rs/vendor/rmcp/src/model/extension.rs
Vendored
+337
View File
@@ -0,0 +1,337 @@
//! A container for those extra data could be carried on request or notification.
//!
//! This file is copied and modified from crate [http](https://github.com/hyperium/http).
//!
//! - Original code license: <https://github.com/hyperium/http/blob/master/LICENSE-MIT>
//! - Original code: <https://github.com/hyperium/http/blob/master/src/extensions.rs>
use std::{
any::{Any, TypeId},
collections::HashMap,
fmt,
hash::{BuildHasherDefault, Hasher},
};
type AnyMap = HashMap<TypeId, Box<dyn AnyClone + Send + Sync>, BuildHasherDefault<IdHasher>>;
// With TypeIds as keys, there's no need to hash them. They are already hashes
// themselves, coming from the compiler. The IdHasher just holds the u64 of
// the TypeId, and then returns it, instead of doing any bit fiddling.
#[derive(Default)]
struct IdHasher(u64);
impl Hasher for IdHasher {
fn write(&mut self, _: &[u8]) {
unreachable!("TypeId calls write_u64");
}
#[inline]
fn write_u64(&mut self, id: u64) {
self.0 = id;
}
#[inline]
fn finish(&self) -> u64 {
self.0
}
}
/// A type map of protocol extensions.
///
/// `Extensions` can be used by `Request` `Notification` and `Response` to store
/// extra data derived from the underlying protocol.
#[derive(Clone, Default)]
pub struct Extensions {
// If extensions are never used, no need to carry around an empty HashMap.
// That's 3 words. Instead, this is only 1 word.
map: Option<Box<AnyMap>>,
}
impl Extensions {
/// Create an empty `Extensions`.
#[inline]
pub const fn new() -> Extensions {
Extensions { map: None }
}
/// Insert a type into this `Extensions`.
///
/// If a extension of this type already existed, it will
/// be returned and replaced with the new one.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// assert!(ext.insert(5i32).is_none());
/// assert!(ext.insert(4u8).is_none());
/// assert_eq!(ext.insert(9i32), Some(5i32));
/// ```
pub fn insert<T: Clone + Send + Sync + 'static>(&mut self, val: T) -> Option<T> {
self.map
.get_or_insert_with(Box::default)
.insert(TypeId::of::<T>(), Box::new(val))
.and_then(|boxed| boxed.into_any().downcast().ok().map(|boxed| *boxed))
}
/// Get a reference to a type previously inserted on this `Extensions`.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// assert!(ext.get::<i32>().is_none());
/// ext.insert(5i32);
///
/// assert_eq!(ext.get::<i32>(), Some(&5i32));
/// ```
pub fn get<T: Send + Sync + 'static>(&self) -> Option<&T> {
self.map
.as_ref()
.and_then(|map| map.get(&TypeId::of::<T>()))
.and_then(|boxed| (**boxed).as_any().downcast_ref())
}
/// Get a mutable reference to a type previously inserted on this `Extensions`.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// ext.insert(String::from("Hello"));
/// ext.get_mut::<String>().unwrap().push_str(" World");
///
/// assert_eq!(ext.get::<String>().unwrap(), "Hello World");
/// ```
pub fn get_mut<T: Send + Sync + 'static>(&mut self) -> Option<&mut T> {
self.map
.as_mut()
.and_then(|map| map.get_mut(&TypeId::of::<T>()))
.and_then(|boxed| (**boxed).as_any_mut().downcast_mut())
}
/// Get a mutable reference to a type, inserting `value` if not already present on this
/// `Extensions`.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// *ext.get_or_insert(1i32) += 2;
///
/// assert_eq!(*ext.get::<i32>().unwrap(), 3);
/// ```
pub fn get_or_insert<T: Clone + Send + Sync + 'static>(&mut self, value: T) -> &mut T {
self.get_or_insert_with(|| value)
}
/// Get a mutable reference to a type, inserting the value created by `f` if not already present
/// on this `Extensions`.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// *ext.get_or_insert_with(|| 1i32) += 2;
///
/// assert_eq!(*ext.get::<i32>().unwrap(), 3);
/// ```
pub fn get_or_insert_with<T: Clone + Send + Sync + 'static, F: FnOnce() -> T>(
&mut self,
f: F,
) -> &mut T {
let out = self
.map
.get_or_insert_with(Box::default)
.entry(TypeId::of::<T>())
.or_insert_with(|| Box::new(f()));
(**out).as_any_mut().downcast_mut().unwrap()
}
/// Get a mutable reference to a type, inserting the type's default value if not already present
/// on this `Extensions`.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// *ext.get_or_insert_default::<i32>() += 2;
///
/// assert_eq!(*ext.get::<i32>().unwrap(), 2);
/// ```
pub fn get_or_insert_default<T: Default + Clone + Send + Sync + 'static>(&mut self) -> &mut T {
self.get_or_insert_with(T::default)
}
/// Remove a type from this `Extensions`.
///
/// If a extension of this type existed, it will be returned.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// ext.insert(5i32);
/// assert_eq!(ext.remove::<i32>(), Some(5i32));
/// assert!(ext.get::<i32>().is_none());
/// ```
pub fn remove<T: Send + Sync + 'static>(&mut self) -> Option<T> {
self.map
.as_mut()
.and_then(|map| map.remove(&TypeId::of::<T>()))
.and_then(|boxed| boxed.into_any().downcast().ok().map(|boxed| *boxed))
}
/// Clear the `Extensions` of all inserted extensions.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// ext.insert(5i32);
/// ext.clear();
///
/// assert!(ext.get::<i32>().is_none());
/// ```
#[inline]
pub fn clear(&mut self) {
if let Some(ref mut map) = self.map {
map.clear();
}
}
/// Check whether the extension set is empty or not.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// assert!(ext.is_empty());
/// ext.insert(5i32);
/// assert!(!ext.is_empty());
/// ```
#[inline]
pub fn is_empty(&self) -> bool {
self.map.as_ref().is_none_or(|map| map.is_empty())
}
/// Get the number of extensions available.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext = Extensions::new();
/// assert_eq!(ext.len(), 0);
/// ext.insert(5i32);
/// assert_eq!(ext.len(), 1);
/// ```
#[inline]
pub fn len(&self) -> usize {
self.map.as_ref().map_or(0, |map| map.len())
}
/// Extends `self` with another `Extensions`.
///
/// If an instance of a specific type exists in both, the one in `self` is overwritten with the
/// one from `other`.
///
/// # Example
///
/// ```
/// # use rmcp::model::Extensions;
/// let mut ext_a = Extensions::new();
/// ext_a.insert(8u8);
/// ext_a.insert(16u16);
///
/// let mut ext_b = Extensions::new();
/// ext_b.insert(4u8);
/// ext_b.insert("hello");
///
/// ext_a.extend(ext_b);
/// assert_eq!(ext_a.len(), 3);
/// assert_eq!(ext_a.get::<u8>(), Some(&4u8));
/// assert_eq!(ext_a.get::<u16>(), Some(&16u16));
/// assert_eq!(ext_a.get::<&'static str>().copied(), Some("hello"));
/// ```
pub fn extend(&mut self, other: Self) {
if let Some(other) = other.map {
if let Some(map) = &mut self.map {
map.extend(*other);
} else {
self.map = Some(other);
}
}
}
}
impl fmt::Debug for Extensions {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Extensions").finish()
}
}
trait AnyClone: Any {
fn clone_box(&self) -> Box<dyn AnyClone + Send + Sync>;
fn as_any(&self) -> &dyn Any;
fn as_any_mut(&mut self) -> &mut dyn Any;
fn into_any(self: Box<Self>) -> Box<dyn Any>;
}
impl<T: Clone + Send + Sync + 'static> AnyClone for T {
fn clone_box(&self) -> Box<dyn AnyClone + Send + Sync> {
Box::new(self.clone())
}
fn as_any(&self) -> &dyn Any {
self
}
fn as_any_mut(&mut self) -> &mut dyn Any {
self
}
fn into_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}
impl Clone for Box<dyn AnyClone + Send + Sync> {
fn clone(&self) -> Self {
(**self).clone_box()
}
}
#[test]
fn test_extensions() {
#[derive(Clone, Debug, PartialEq)]
struct MyType(i32);
let mut extensions = Extensions::new();
extensions.insert(5i32);
extensions.insert(MyType(10));
assert_eq!(extensions.get(), Some(&5i32));
assert_eq!(extensions.get_mut(), Some(&mut 5i32));
let ext2 = extensions.clone();
assert_eq!(extensions.remove::<i32>(), Some(5i32));
assert!(extensions.get::<i32>().is_none());
// clone still has it
assert_eq!(ext2.get(), Some(&5i32));
assert_eq!(ext2.get(), Some(&MyType(10)));
assert_eq!(extensions.get::<bool>(), None);
assert_eq!(extensions.get(), Some(&MyType(10)));
}
codex-rs/vendor/rmcp/src/model/meta.rs
Vendored
+289
View File
@@ -0,0 +1,289 @@
use std::ops::{Deref, DerefMut};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::{
ClientNotification, ClientRequest, CustomNotification, CustomRequest, Extensions, JsonObject,
JsonRpcMessage, NumberOrString, ProgressToken, ServerNotification, ServerRequest,
};
pub trait GetMeta {
fn get_meta_mut(&mut self) -> &mut Meta;
fn get_meta(&self) -> &Meta;
}
pub trait GetExtensions {
fn extensions(&self) -> &Extensions;
fn extensions_mut(&mut self) -> &mut Extensions;
}
/// Trait for request params that contain the `_meta` field.
///
/// Per the MCP 2025-11-25 spec, all request params should have an optional `_meta`
/// field that can contain a `progressToken` for tracking long-running operations.
pub trait RequestParamsMeta {
/// Get a reference to the meta field
fn meta(&self) -> Option<&Meta>;
/// Get a mutable reference to the meta field
fn meta_mut(&mut self) -> &mut Option<Meta>;
/// Set the meta field
fn set_meta(&mut self, meta: Meta) {
*self.meta_mut() = Some(meta);
}
/// Get the progress token from meta, if present
fn progress_token(&self) -> Option<ProgressToken> {
self.meta().and_then(|m| m.get_progress_token())
}
/// Set a progress token in meta
fn set_progress_token(&mut self, token: ProgressToken) {
match self.meta_mut() {
Some(meta) => meta.set_progress_token(token),
none => {
let mut meta = Meta::new();
meta.set_progress_token(token);
*none = Some(meta);
}
}
}
}
/// Trait for task-augmented request params that contain both `_meta` and `task` fields.
///
/// Per the MCP 2025-11-25 spec, certain requests (like `tools/call` and `sampling/createMessage`)
/// can include a `task` field to signal that the caller wants task-augmented execution.
pub trait TaskAugmentedRequestParamsMeta: RequestParamsMeta {
/// Get a reference to the task field
fn task(&self) -> Option<&JsonObject>;
/// Get a mutable reference to the task field
fn task_mut(&mut self) -> &mut Option<JsonObject>;
/// Set the task field
fn set_task(&mut self, task: JsonObject) {
*self.task_mut() = Some(task);
}
}
impl GetExtensions for CustomNotification {
fn extensions(&self) -> &Extensions {
&self.extensions
}
fn extensions_mut(&mut self) -> &mut Extensions {
&mut self.extensions
}
}
impl GetMeta for CustomNotification {
fn get_meta_mut(&mut self) -> &mut Meta {
self.extensions_mut().get_or_insert_default()
}
fn get_meta(&self) -> &Meta {
self.extensions()
.get::<Meta>()
.unwrap_or(Meta::static_empty())
}
}
impl GetExtensions for CustomRequest {
fn extensions(&self) -> &Extensions {
&self.extensions
}
fn extensions_mut(&mut self) -> &mut Extensions {
&mut self.extensions
}
}
impl GetMeta for CustomRequest {
fn get_meta_mut(&mut self) -> &mut Meta {
self.extensions_mut().get_or_insert_default()
}
fn get_meta(&self) -> &Meta {
self.extensions()
.get::<Meta>()
.unwrap_or(Meta::static_empty())
}
}
macro_rules! variant_extension {
(
$Enum: ident {
$($variant: ident)*
}
) => {
impl GetExtensions for $Enum {
fn extensions(&self) -> &Extensions {
match self {
$(
$Enum::$variant(v) => &v.extensions,
)*
}
}
fn extensions_mut(&mut self) -> &mut Extensions {
match self {
$(
$Enum::$variant(v) => &mut v.extensions,
)*
}
}
}
impl GetMeta for $Enum {
fn get_meta_mut(&mut self) -> &mut Meta {
self.extensions_mut().get_or_insert_default()
}
fn get_meta(&self) -> &Meta {
self.extensions().get::<Meta>().unwrap_or(Meta::static_empty())
}
}
};
}
variant_extension! {
ClientRequest {
PingRequest
InitializeRequest
CompleteRequest
SetLevelRequest
GetPromptRequest
ListPromptsRequest
ListResourcesRequest
ListResourceTemplatesRequest
ReadResourceRequest
SubscribeRequest
UnsubscribeRequest
CallToolRequest
ListToolsRequest
CustomRequest
GetTaskInfoRequest
ListTasksRequest
GetTaskResultRequest
CancelTaskRequest
}
}
variant_extension! {
ServerRequest {
PingRequest
CreateMessageRequest
ListRootsRequest
CreateElicitationRequest
CustomRequest
}
}
variant_extension! {
ClientNotification {
CancelledNotification
ProgressNotification
InitializedNotification
RootsListChangedNotification
CustomNotification
}
}
variant_extension! {
ServerNotification {
CancelledNotification
ProgressNotification
LoggingMessageNotification
ResourceUpdatedNotification
ResourceListChangedNotification
ToolListChangedNotification
PromptListChangedNotification
ElicitationCompletionNotification
CustomNotification
}
}
#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
#[serde(transparent)]
pub struct Meta(pub JsonObject);
const PROGRESS_TOKEN_FIELD: &str = "progressToken";
impl Meta {
pub fn new() -> Self {
Self(JsonObject::new())
}
/// Create a new Meta with a progress token set
pub fn with_progress_token(token: ProgressToken) -> Self {
let mut meta = Self::new();
meta.set_progress_token(token);
meta
}
pub(crate) fn static_empty() -> &'static Self {
static EMPTY: std::sync::OnceLock<Meta> = std::sync::OnceLock::new();
EMPTY.get_or_init(Default::default)
}
pub fn get_progress_token(&self) -> Option<ProgressToken> {
self.0.get(PROGRESS_TOKEN_FIELD).and_then(|v| match v {
Value::String(s) => Some(ProgressToken(NumberOrString::String(s.to_string().into()))),
Value::Number(n) => {
if let Some(i) = n.as_i64() {
Some(ProgressToken(NumberOrString::Number(i)))
} else if let Some(u) = n.as_u64() {
if u <= i64::MAX as u64 {
Some(ProgressToken(NumberOrString::Number(u as i64)))
} else {
None
}
} else {
None
}
}
_ => None,
})
}
pub fn set_progress_token(&mut self, token: ProgressToken) {
match token.0 {
NumberOrString::String(ref s) => self.0.insert(
PROGRESS_TOKEN_FIELD.to_string(),
Value::String(s.to_string()),
),
NumberOrString::Number(n) => self
.0
.insert(PROGRESS_TOKEN_FIELD.to_string(), Value::Number(n.into())),
};
}
pub fn extend(&mut self, other: Meta) {
for (k, v) in other.0.into_iter() {
self.0.insert(k, v);
}
}
}
impl Deref for Meta {
type Target = JsonObject;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl DerefMut for Meta {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl<Req, Resp, Noti> JsonRpcMessage<Req, Resp, Noti>
where
Req: GetExtensions,
Noti: GetExtensions,
{
pub fn insert_extension<T: Clone + Send + Sync + 'static>(&mut self, value: T) {
match self {
JsonRpcMessage::Request(json_rpc_request) => {
json_rpc_request.request.extensions_mut().insert(value);
}
JsonRpcMessage::Notification(json_rpc_notification) => {
json_rpc_notification
.notification
.extensions_mut()
.insert(value);
}
_ => {}
}
}
}
codex-rs/vendor/rmcp/src/model/prompt.rs
Vendored
+271
View File
@@ -0,0 +1,271 @@
use serde::{Deserialize, Serialize};
use super::{
AnnotateAble, Annotations, Icon, Meta, RawEmbeddedResource,
content::{EmbeddedResource, ImageContent},
resource::ResourceContents,
};
/// A prompt that can be used to generate text from a model
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Prompt {
/// The name of the prompt
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<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>>,
/// Optional list of icons for the prompt
#[serde(skip_serializing_if = "Option::is_none")]
pub icons: Option<Vec<Icon>>,
/// Optional additional metadata for this prompt
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<Meta>,
}
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(),
title: None,
description: description.map(Into::into),
arguments,
icons: None,
meta: None,
}
}
}
/// Represents a prompt argument that can be passed to customize the prompt
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct PromptArgument {
/// The name of the argument
pub name: String,
/// A human-readable title for the argument
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<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")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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 = "snake_case")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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 },
/// A link to a resource that can be fetched separately
ResourceLink {
#[serde(flatten)]
link: super::resource::Resource,
},
}
impl PromptMessageContent {
pub fn text(text: impl Into<String>) -> Self {
Self::Text { text: text.into() }
}
/// Create a resource link content
pub fn resource_link(resource: super::resource::Resource) -> Self {
Self::ResourceLink { link: resource }
}
}
/// A message in a prompt conversation
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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() },
}
}
/// Create a new image message. `meta` and `annotations` are optional.
#[cfg(feature = "base64")]
pub fn new_image(
role: PromptMessageRole,
data: &[u8],
mime_type: &str,
meta: Option<crate::model::Meta>,
annotations: Option<Annotations>,
) -> Self {
use base64::{Engine, prelude::BASE64_STANDARD};
let base64 = BASE64_STANDARD.encode(data);
Self {
role,
content: PromptMessageContent::Image {
image: crate::model::RawImageContent {
data: base64,
mime_type: mime_type.into(),
meta,
}
.optional_annotate(annotations),
},
}
}
/// Create a new resource message. `resource_meta`, `resource_content_meta`, and `annotations` are optional.
pub fn new_resource(
role: PromptMessageRole,
uri: String,
mime_type: Option<String>,
text: Option<String>,
resource_meta: Option<crate::model::Meta>,
resource_content_meta: Option<crate::model::Meta>,
annotations: Option<Annotations>,
) -> Self {
let resource_contents = match text {
Some(t) => ResourceContents::TextResourceContents {
uri,
mime_type,
text: t,
meta: resource_content_meta,
},
None => ResourceContents::BlobResourceContents {
uri,
mime_type,
blob: String::new(),
meta: resource_content_meta,
},
};
Self {
role,
content: PromptMessageContent::Resource {
resource: RawEmbeddedResource {
meta: resource_meta,
resource: resource_contents,
}
.optional_annotate(annotations),
},
}
}
/// Note: PromptMessage text content does not carry protocol-level _meta per current schema.
/// This function exists for API symmetry but ignores the meta parameter.
pub fn new_text_with_meta<S: Into<String>>(
role: PromptMessageRole,
text: S,
_meta: Option<crate::model::Meta>,
) -> Self {
Self::new_text(role, text)
}
/// Create a new resource link message
pub fn new_resource_link(role: PromptMessageRole, resource: super::resource::Resource) -> Self {
Self {
role,
content: PromptMessageContent::ResourceLink { link: resource },
}
}
}
#[cfg(test)]
mod tests {
use serde_json;
use super::*;
#[test]
fn test_prompt_message_image_serialization() {
let image_content = crate::model::RawImageContent {
data: "base64data".to_string(),
mime_type: "image/png".to_string(),
meta: None,
};
let json = serde_json::to_string(&image_content).unwrap();
println!("PromptMessage ImageContent JSON: {}", json);
// Verify it contains mimeType (camelCase) not mime_type (snake_case)
assert!(json.contains("mimeType"));
assert!(!json.contains("mime_type"));
}
#[test]
fn test_prompt_message_resource_link_serialization() {
use super::super::resource::RawResource;
let resource = RawResource::new("file:///test.txt", "test.txt");
let message =
PromptMessage::new_resource_link(PromptMessageRole::User, resource.no_annotation());
let json = serde_json::to_string(&message).unwrap();
println!("PromptMessage with ResourceLink JSON: {}", json);
// Verify it contains the correct type tag
assert!(json.contains("\"type\":\"resource_link\""));
assert!(json.contains("\"uri\":\"file:///test.txt\""));
assert!(json.contains("\"name\":\"test.txt\""));
}
#[test]
fn test_prompt_message_content_resource_link_deserialization() {
let json = r#"{
"type": "resource_link",
"uri": "file:///example.txt",
"name": "example.txt",
"description": "Example file",
"mimeType": "text/plain"
}"#;
let content: PromptMessageContent = serde_json::from_str(json).unwrap();
if let PromptMessageContent::ResourceLink { link } = content {
assert_eq!(link.uri, "file:///example.txt");
assert_eq!(link.name, "example.txt");
assert_eq!(link.description, Some("Example file".to_string()));
assert_eq!(link.mime_type, Some("text/plain".to_string()));
} else {
panic!("Expected ResourceLink variant");
}
}
}
codex-rs/vendor/rmcp/src/model/resource.rs
Vendored
+188
View File
@@ -0,0 +1,188 @@
use serde::{Deserialize, Serialize};
use super::{Annotated, Icon, Meta};
/// Represents a resource in the extension with metadata
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
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,
/// Human-readable title of the resource
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<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>,
/// Optional list of icons for the resource
#[serde(skip_serializing_if = "Option::is_none")]
pub icons: Option<Vec<Icon>>,
/// Optional additional metadata for this resource
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<Meta>,
}
pub type Resource = Annotated<RawResource>;
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct RawResourceTemplate {
pub uri_template: String,
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub mime_type: Option<String>,
/// Optional list of icons for the resource template
#[serde(skip_serializing_if = "Option::is_none")]
pub icons: Option<Vec<Icon>>,
}
pub type ResourceTemplate = Annotated<RawResourceTemplate>;
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
#[serde(untagged)]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub enum ResourceContents {
#[serde(rename_all = "camelCase")]
TextResourceContents {
uri: String,
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
text: String,
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
meta: Option<Meta>,
},
#[serde(rename_all = "camelCase")]
BlobResourceContents {
uri: String,
#[serde(skip_serializing_if = "Option::is_none")]
mime_type: Option<String>,
blob: String,
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
meta: Option<Meta>,
},
}
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(),
meta: None,
}
}
}
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(),
title: None,
description: None,
mime_type: None,
size: None,
icons: None,
meta: None,
}
}
}
#[cfg(test)]
mod tests {
use serde_json;
use super::*;
#[test]
fn test_resource_serialization() {
let resource = RawResource {
uri: "file:///test.txt".to_string(),
title: None,
name: "test".to_string(),
description: Some("Test resource".to_string()),
mime_type: Some("text/plain".to_string()),
size: Some(100),
icons: None,
meta: None,
};
let json = serde_json::to_string(&resource).unwrap();
println!("Serialized JSON: {}", json);
// Verify it contains mimeType (camelCase) not mime_type (snake_case)
assert!(json.contains("mimeType"));
assert!(!json.contains("mime_type"));
}
#[test]
fn test_resource_contents_serialization() {
let text_contents = ResourceContents::TextResourceContents {
uri: "file:///test.txt".to_string(),
mime_type: Some("text/plain".to_string()),
text: "Hello world".to_string(),
meta: None,
};
let json = serde_json::to_string(&text_contents).unwrap();
println!("ResourceContents JSON: {}", json);
// Verify it contains mimeType (camelCase) not mime_type (snake_case)
assert!(json.contains("mimeType"));
assert!(!json.contains("mime_type"));
}
#[test]
fn test_resource_template_with_icons() {
let resource_template = RawResourceTemplate {
uri_template: "file:///{path}".to_string(),
name: "template".to_string(),
title: Some("Test Template".to_string()),
description: Some("A test resource template".to_string()),
mime_type: Some("text/plain".to_string()),
icons: Some(vec![Icon {
src: "https://example.com/icon.png".to_string(),
mime_type: Some("image/png".to_string()),
sizes: Some(vec!["48x48".to_string()]),
}]),
};
let json = serde_json::to_value(&resource_template).unwrap();
assert!(json["icons"].is_array());
assert_eq!(json["icons"][0]["src"], "https://example.com/icon.png");
assert_eq!(json["icons"][0]["sizes"][0], "48x48");
}
#[test]
fn test_resource_template_without_icons() {
let resource_template = RawResourceTemplate {
uri_template: "file:///{path}".to_string(),
name: "template".to_string(),
title: None,
description: None,
mime_type: None,
icons: None,
};
let json = serde_json::to_value(&resource_template).unwrap();
assert!(json.get("icons").is_none());
}
}
codex-rs/vendor/rmcp/src/model/serde_impl.rs
Vendored
+373
View File
@@ -0,0 +1,373 @@
use std::borrow::Cow;
use serde::{Deserialize, Serialize};
use super::{
CustomNotification, CustomRequest, Extensions, Meta, Notification, NotificationNoParam,
Request, RequestNoParam, RequestOptionalParam,
};
#[derive(Serialize, Deserialize)]
struct WithMeta<'a, P> {
#[serde(skip_serializing_if = "Option::is_none")]
_meta: Option<Cow<'a, Meta>>,
#[serde(flatten)]
_rest: P,
}
#[derive(Serialize, Deserialize)]
struct Proxy<'a, M, P> {
method: M,
params: WithMeta<'a, P>,
}
#[derive(Serialize, Deserialize)]
struct ProxyOptionalParam<'a, M, P> {
method: M,
params: Option<WithMeta<'a, P>>,
}
#[derive(Serialize, Deserialize)]
struct ProxyNoParam<M> {
method: M,
}
impl<M, R> Serialize for Request<M, R>
where
M: Serialize,
R: Serialize,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let extensions = &self.extensions;
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
Proxy::serialize(
&Proxy {
method: &self.method,
params: WithMeta {
_rest: &self.params,
_meta,
},
},
serializer,
)
}
}
impl<'de, M, R> Deserialize<'de> for Request<M, R>
where
M: Deserialize<'de>,
R: Deserialize<'de>,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let body = Proxy::deserialize(deserializer)?;
let _meta = body.params._meta.map(|m| m.into_owned());
let mut extensions = Extensions::new();
if let Some(meta) = _meta {
extensions.insert(meta);
}
Ok(Request {
extensions,
method: body.method,
params: body.params._rest,
})
}
}
impl<M, R> Serialize for RequestOptionalParam<M, R>
where
M: Serialize,
R: Serialize,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let extensions = &self.extensions;
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
Proxy::serialize(
&Proxy {
method: &self.method,
params: WithMeta {
_rest: &self.params,
_meta,
},
},
serializer,
)
}
}
impl<'de, M, R> Deserialize<'de> for RequestOptionalParam<M, R>
where
M: Deserialize<'de>,
R: Deserialize<'de>,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let body = ProxyOptionalParam::<'_, _, Option<R>>::deserialize(deserializer)?;
let mut params = None;
let mut _meta = None;
if let Some(body_params) = body.params {
params = body_params._rest;
_meta = body_params._meta.map(|m| m.into_owned());
}
let mut extensions = Extensions::new();
if let Some(meta) = _meta {
extensions.insert(meta);
}
Ok(RequestOptionalParam {
extensions,
method: body.method,
params,
})
}
}
impl<M> Serialize for RequestNoParam<M>
where
M: Serialize,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let extensions = &self.extensions;
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
ProxyNoParam::serialize(
&ProxyNoParam {
method: &self.method,
},
serializer,
)
}
}
impl<'de, M> Deserialize<'de> for RequestNoParam<M>
where
M: Deserialize<'de>,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let body = ProxyNoParam::<_>::deserialize(deserializer)?;
let extensions = Extensions::new();
Ok(RequestNoParam {
extensions,
method: body.method,
})
}
}
impl<M, R> Serialize for Notification<M, R>
where
M: Serialize,
R: Serialize,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let extensions = &self.extensions;
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
Proxy::serialize(
&Proxy {
method: &self.method,
params: WithMeta {
_rest: &self.params,
_meta,
},
},
serializer,
)
}
}
impl<'de, M, R> Deserialize<'de> for Notification<M, R>
where
M: Deserialize<'de>,
R: Deserialize<'de>,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let body = Proxy::deserialize(deserializer)?;
let _meta = body.params._meta.map(|m| m.into_owned());
let mut extensions = Extensions::new();
if let Some(meta) = _meta {
extensions.insert(meta);
}
Ok(Notification {
extensions,
method: body.method,
params: body.params._rest,
})
}
}
impl<M> Serialize for NotificationNoParam<M>
where
M: Serialize,
{
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let extensions = &self.extensions;
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
ProxyNoParam::serialize(
&ProxyNoParam {
method: &self.method,
},
serializer,
)
}
}
impl<'de, M> Deserialize<'de> for NotificationNoParam<M>
where
M: Deserialize<'de>,
{
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let body = ProxyNoParam::<_>::deserialize(deserializer)?;
let extensions = Extensions::new();
Ok(NotificationNoParam {
extensions,
method: body.method,
})
}
}
impl Serialize for CustomRequest {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let extensions = &self.extensions;
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
let params = self.params.as_ref();
let params = if _meta.is_some() || params.is_some() {
Some(WithMeta {
_meta,
_rest: &self.params,
})
} else {
None
};
ProxyOptionalParam::serialize(
&ProxyOptionalParam {
method: &self.method,
params,
},
serializer,
)
}
}
impl<'de> Deserialize<'de> for CustomRequest {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let body =
ProxyOptionalParam::<'_, _, Option<serde_json::Value>>::deserialize(deserializer)?;
let mut params = None;
let mut _meta = None;
if let Some(body_params) = body.params {
params = body_params._rest;
_meta = body_params._meta.map(|m| m.into_owned());
}
let mut extensions = Extensions::new();
if let Some(meta) = _meta {
extensions.insert(meta);
}
Ok(CustomRequest {
extensions,
method: body.method,
params,
})
}
}
impl Serialize for CustomNotification {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let extensions = &self.extensions;
let _meta = extensions.get::<Meta>().map(Cow::Borrowed);
let params = self.params.as_ref();
let params = if _meta.is_some() || params.is_some() {
Some(WithMeta {
_meta,
_rest: &self.params,
})
} else {
None
};
ProxyOptionalParam::serialize(
&ProxyOptionalParam {
method: &self.method,
params,
},
serializer,
)
}
}
impl<'de> Deserialize<'de> for CustomNotification {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let body =
ProxyOptionalParam::<'_, _, Option<serde_json::Value>>::deserialize(deserializer)?;
let mut params = None;
let mut _meta = None;
if let Some(body_params) = body.params {
params = body_params._rest;
_meta = body_params._meta.map(|m| m.into_owned());
}
let mut extensions = Extensions::new();
if let Some(meta) = _meta {
extensions.insert(meta);
}
Ok(CustomNotification {
extensions,
method: body.method,
params,
})
}
}
#[cfg(test)]
mod test {
use serde_json::json;
use crate::model::ListToolsRequest;
#[test]
fn test_deserialize_lost_tools_request() {
let _req: ListToolsRequest = serde_json::from_value(json!(
{
"method": "tools/list",
}
))
.unwrap();
}
}
codex-rs/vendor/rmcp/src/model/task.rs
Vendored
+79
View File
@@ -0,0 +1,79 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// Canonical task lifecycle status as defined by SEP-1686.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub enum TaskStatus {
/// The receiver accepted the request and is currently working on it.
#[default]
Working,
/// The receiver requires additional input before work can continue.
InputRequired,
/// The underlying operation completed successfully and the result is ready.
Completed,
/// The underlying operation failed and will not continue.
Failed,
/// The task was cancelled and will not continue processing.
Cancelled,
}
/// Final result for a succeeded task (returned from `tasks/result`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct TaskResult {
/// MIME type or custom content-type identifier.
pub content_type: String,
/// The actual result payload, matching the underlying request's schema.
pub value: Value,
/// Optional short summary for UI surfaces.
#[serde(skip_serializing_if = "Option::is_none")]
pub summary: Option<String>,
}
/// Primary Task object that surfaces metadata during the task lifecycle.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Task {
/// Unique task identifier generated by the receiver.
pub task_id: String,
/// Current lifecycle status (see [`TaskStatus`]).
pub status: TaskStatus,
/// Optional human-readable status message for UI surfaces.
#[serde(skip_serializing_if = "Option::is_none")]
pub status_message: Option<String>,
/// ISO-8601 creation timestamp.
pub created_at: String,
/// ISO-8601 timestamp for the most recent status change.
#[serde(skip_serializing_if = "Option::is_none")]
pub last_updated_at: Option<String>,
/// Retention window in milliseconds that the receiver agreed to honor.
#[serde(skip_serializing_if = "Option::is_none")]
pub ttl: Option<u64>,
/// Suggested polling interval (milliseconds).
#[serde(skip_serializing_if = "Option::is_none")]
pub poll_interval: Option<u64>,
}
/// Wrapper returned by task-augmented requests (CreateTaskResult in SEP-1686).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct CreateTaskResult {
pub task: Task,
}
/// Paginated list of tasks
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct TaskList {
pub tasks: Vec<Task>,
#[serde(skip_serializing_if = "Option::is_none")]
pub next_cursor: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub total: Option<u64>,
}
codex-rs/vendor/rmcp/src/model/tool.rs
Vendored
+260
View File
@@ -0,0 +1,260 @@
use std::{borrow::Cow, sync::Arc};
use schemars::JsonSchema;
/// 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::{Icon, JsonObject, Meta};
/// A tool that can be used by a model.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct Tool {
/// The name of the tool
pub name: Cow<'static, str>,
/// A human-readable title for the tool
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
/// A description of what the tool does
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<Cow<'static, str>>,
/// A JSON Schema object defining the expected parameters for the tool
pub input_schema: Arc<JsonObject>,
/// An optional JSON Schema object defining the structure of the tool's output
#[serde(skip_serializing_if = "Option::is_none")]
pub output_schema: Option<Arc<JsonObject>>,
#[serde(skip_serializing_if = "Option::is_none")]
/// Optional additional tool information.
pub annotations: Option<ToolAnnotations>,
/// Execution-related configuration including task support mode.
#[serde(skip_serializing_if = "Option::is_none")]
pub execution: Option<ToolExecution>,
/// Optional list of icons for the tool
#[serde(skip_serializing_if = "Option::is_none")]
pub icons: Option<Vec<Icon>>,
/// Optional additional metadata for this tool
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
pub meta: Option<Meta>,
}
/// Per-tool task support mode as defined in the MCP specification.
///
/// This enum indicates whether a tool supports task-based invocation,
/// allowing clients to know how to properly call the tool.
///
/// See [Tool-Level Negotiation](https://modelcontextprotocol.io/specification/2025-11-25/basic/utilities/tasks#tool-level-negotiation).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "lowercase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub enum TaskSupport {
/// Clients MUST NOT invoke this tool as a task (default behavior).
#[default]
Forbidden,
/// Clients MAY invoke this tool as either a task or a normal call.
Optional,
/// Clients MUST invoke this tool as a task.
Required,
}
/// Execution-related configuration for a tool.
///
/// This struct contains settings that control how a tool should be executed,
/// including task support configuration.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ToolExecution {
/// Indicates whether this tool supports task-based invocation.
///
/// When not present or set to `Forbidden`, clients MUST NOT invoke this tool as a task.
/// When set to `Optional`, clients MAY invoke this tool as a task or normal call.
/// When set to `Required`, clients MUST invoke this tool as a task.
#[serde(skip_serializing_if = "Option::is_none")]
pub task_support: Option<TaskSupport>,
}
impl ToolExecution {
/// Create a new empty ToolExecution configuration.
pub fn new() -> Self {
Self::default()
}
/// Set the task support mode.
pub fn with_task_support(mut self, task_support: TaskSupport) -> Self {
self.task_support = Some(task_support);
self
}
}
/// Additional properties describing a Tool to clients.
///
/// NOTE: all properties in ToolAnnotations are **hints**.
/// They are not guaranteed to provide a faithful description of
/// tool behavior (including descriptive properties like `title`).
///
/// Clients should never make tool use decisions based on ToolAnnotations
/// received from untrusted servers.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
pub struct ToolAnnotations {
/// A human-readable title for the tool.
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
/// If true, the tool does not modify its environment.
///
/// Default: false
#[serde(skip_serializing_if = "Option::is_none")]
pub read_only_hint: Option<bool>,
/// If true, the tool may perform destructive updates to its environment.
/// If false, the tool performs only additive updates.
///
/// (This property is meaningful only when `readOnlyHint == false`)
///
/// Default: true
/// A human-readable description of the tool's purpose.
#[serde(skip_serializing_if = "Option::is_none")]
pub destructive_hint: Option<bool>,
/// If true, calling the tool repeatedly with the same arguments
/// will have no additional effect on the its environment.
///
/// (This property is meaningful only when `readOnlyHint == false`)
///
/// Default: false.
#[serde(skip_serializing_if = "Option::is_none")]
pub idempotent_hint: Option<bool>,
/// If true, this tool may interact with an "open world" of external
/// entities. If false, the tool's domain of interaction is closed.
/// For example, the world of a web search tool is open, whereas that
/// of a memory tool is not.
///
/// Default: true
#[serde(skip_serializing_if = "Option::is_none")]
pub open_world_hint: Option<bool>,
}
impl ToolAnnotations {
pub fn new() -> Self {
Self::default()
}
pub fn with_title<T>(title: T) -> Self
where
T: Into<String>,
{
ToolAnnotations {
title: Some(title.into()),
..Self::default()
}
}
pub fn read_only(self, read_only: bool) -> Self {
ToolAnnotations {
read_only_hint: Some(read_only),
..self
}
}
pub fn destructive(self, destructive: bool) -> Self {
ToolAnnotations {
destructive_hint: Some(destructive),
..self
}
}
pub fn idempotent(self, idempotent: bool) -> Self {
ToolAnnotations {
idempotent_hint: Some(idempotent),
..self
}
}
pub fn open_world(self, open_world: bool) -> Self {
ToolAnnotations {
open_world_hint: Some(open_world),
..self
}
}
/// If not set, defaults to true.
pub fn is_destructive(&self) -> bool {
self.destructive_hint.unwrap_or(true)
}
/// If not set, defaults to false.
pub fn is_idempotent(&self) -> bool {
self.idempotent_hint.unwrap_or(false)
}
}
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(),
title: None,
description: Some(description.into()),
input_schema: input_schema.into(),
output_schema: None,
annotations: None,
execution: None,
icons: None,
meta: None,
}
}
pub fn annotate(self, annotations: ToolAnnotations) -> Self {
Tool {
annotations: Some(annotations),
..self
}
}
/// Set the execution configuration for this tool.
pub fn with_execution(self, execution: ToolExecution) -> Self {
Tool {
execution: Some(execution),
..self
}
}
/// Returns the task support mode for this tool.
///
/// Returns `TaskSupport::Forbidden` if not explicitly set.
pub fn task_support(&self) -> TaskSupport {
self.execution
.as_ref()
.and_then(|e| e.task_support)
.unwrap_or_default()
}
/// Set the output schema using a type that implements JsonSchema
///
/// # Panics
///
/// Panics if the generated schema does not have root type "object" as required by MCP specification.
pub fn with_output_schema<T: JsonSchema + 'static>(mut self) -> Self {
let schema = crate::handler::server::tool::schema_for_output::<T>()
.unwrap_or_else(|e| panic!("Invalid output schema for tool '{}': {}", self.name, e));
self.output_schema = Some(schema);
self
}
/// Set the input schema using a type that implements JsonSchema
pub fn with_input_schema<T: JsonSchema + 'static>(mut self) -> Self {
self.input_schema = crate::handler::server::tool::schema_for_type::<T>();
self
}
/// Get the schema as json value
pub fn schema_as_json_value(&self) -> Value {
Value::Object(self.input_schema.as_ref().clone())
}
}
codex-rs/vendor/rmcp/src/service.rs
Vendored
+944
View File
@@ -0,0 +1,944 @@
use futures::{FutureExt, future::BoxFuture};
use thiserror::Error;
use crate::{
error::ErrorData as McpError,
model::{
CancelledNotification, CancelledNotificationParam, Extensions, GetExtensions, GetMeta,
JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Meta,
NumberOrString, ProgressToken, RequestId, ServerJsonRpcMessage,
},
transport::{DynamicTransportError, IntoTransport, Transport},
};
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
mod client;
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
pub use client::*;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
mod server;
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
pub use server::*;
#[cfg(feature = "tower")]
#[cfg_attr(docsrs, doc(cfg(feature = "tower")))]
mod tower;
use tokio_util::sync::{CancellationToken, DropGuard};
#[cfg(feature = "tower")]
#[cfg_attr(docsrs, doc(cfg(feature = "tower")))]
pub use tower::*;
use tracing::{Instrument as _, instrument};
#[derive(Error, Debug)]
#[non_exhaustive]
pub enum ServiceError {
#[error("Mcp error: {0}")]
McpError(McpError),
#[error("Transport send error: {0}")]
TransportSend(DynamicTransportError),
#[error("Transport closed")]
TransportClosed,
#[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 },
}
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 + GetMeta + GetExtensions;
type Resp: TransferObject;
type Not: TryInto<CancelledNotification, Error = Self::Not>
+ From<CancelledNotification>
+ TransferObject;
type PeerReq: TransferObject + GetMeta + GetExtensions;
type PeerResp: TransferObject;
type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
+ From<CancelledNotification>
+ TransferObject
+ GetMeta
+ GetExtensions;
type InitializeError;
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 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,
context: NotificationContext<R>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_;
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>, R::InitializeError>> + Send
where
T: IntoTransport<R, E, A>,
E: std::error::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>, R::InitializeError>> + Send
where
T: IntoTransport<R, E, A>,
E: std::error::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,
context: NotificationContext<R>,
) -> impl Future<Output = Result<(), McpError>> + Send + '_ {
DynService::handle_notification(self.as_ref(), notification, context)
}
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,
context: NotificationContext<R>,
) -> BoxFuture<'_, Result<(), McpError>>;
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,
context: NotificationContext<R>,
) -> BoxFuture<'_, Result<(), McpError>> {
Box::pin(self.handle_notification(notification, context))
}
fn get_info(&self) -> R::Info {
self.get_info()
}
}
use std::{
collections::{HashMap, VecDeque},
ops::Deref,
sync::{Arc, atomic::AtomicU64},
time::Duration,
};
use tokio::sync::mpsc;
pub trait RequestIdProvider: Send + Sync + 'static {
fn next_request_id(&self) -> RequestId;
}
pub trait ProgressTokenProvider: Send + Sync + 'static {
fn next_progress_token(&self) -> ProgressToken;
}
pub type AtomicU32RequestIdProvider = AtomicU32Provider;
pub type AtomicU32ProgressTokenProvider = AtomicU32Provider;
#[derive(Debug, Default)]
pub struct AtomicU32Provider {
id: AtomicU64,
}
impl RequestIdProvider for AtomicU32Provider {
fn next_request_id(&self) -> RequestId {
let id = self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
// Safe conversion: we start at 0 and increment by 1, so we won't overflow i64::MAX in practice
RequestId::Number(id as i64)
}
}
impl ProgressTokenProvider for AtomicU32Provider {
fn next_progress_token(&self) -> ProgressToken {
let id = self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
ProgressToken(NumberOrString::Number(id as i64))
}
}
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,
pub progress_token: ProgressToken,
}
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::TransportClosed)?
})
.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,
extensions: Default::default(),
};
let _ = self.peer.send_notification(notification.into()).await;
error
}
}
} else {
self.rx.await.map_err(|_e| ServiceError::TransportClosed)?
}
}
/// 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,
extensions: Default::default(),
};
self.peer.send_notification(notification.into()).await?;
Ok(())
}
}
#[derive(Debug)]
pub(crate) enum PeerSinkMessage<R: ServiceRole> {
Request {
request: R::Req,
id: RequestId,
responder: Responder<Result<R::PeerResp, ServiceError>>,
},
Notification {
notification: R::Not,
responder: 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_request_with_option`].
#[derive(Clone)]
pub struct Peer<R: ServiceRole> {
tx: mpsc::Sender<PeerSinkMessage<R>>,
request_id_provider: Arc<dyn RequestIdProvider>,
progress_token_provider: Arc<dyn ProgressTokenProvider>,
info: Arc<tokio::sync::OnceCell<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 {
pub timeout: Option<Duration>,
pub meta: Option<Meta>,
}
impl PeerRequestOptions {
pub fn no_options() -> Self {
Self::default()
}
}
impl<R: ServiceRole> Peer<R> {
const CLIENT_CHANNEL_BUFFER_SIZE: usize = 1024;
pub(crate) fn new(
request_id_provider: Arc<dyn RequestIdProvider>,
peer_info: Option<R::PeerInfo>,
) -> (Peer<R>, ProxyOutbound<R>) {
let (tx, rx) = mpsc::channel(Self::CLIENT_CHANNEL_BUFFER_SIZE);
(
Self {
tx,
request_id_provider,
progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()),
info: Arc::new(tokio::sync::OnceCell::new_with(peer_info)),
},
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::TransportClosed)?;
receiver.await.map_err(|_e| ServiceError::TransportClosed)?
}
pub async fn send_request(&self, request: R::Req) -> Result<R::PeerResp, ServiceError> {
self.send_request_with_option(request, PeerRequestOptions::no_options())
.await?
.await_response()
.await
}
pub async fn send_cancellable_request(
&self,
request: R::Req,
options: PeerRequestOptions,
) -> Result<RequestHandle<R>, ServiceError> {
self.send_request_with_option(request, options).await
}
pub async fn send_request_with_option(
&self,
mut request: R::Req,
options: PeerRequestOptions,
) -> Result<RequestHandle<R>, ServiceError> {
let id = self.request_id_provider.next_request_id();
let progress_token = self.progress_token_provider.next_progress_token();
request
.get_meta_mut()
.set_progress_token(progress_token.clone());
if let Some(meta) = options.meta.clone() {
request.get_meta_mut().extend(meta);
}
let (responder, receiver) = tokio::sync::oneshot::channel();
self.tx
.send(PeerSinkMessage::Request {
request,
id: id.clone(),
responder,
})
.await
.map_err(|_m| ServiceError::TransportClosed)?;
Ok(RequestHandle {
id,
rx: receiver,
progress_token,
options,
peer: self.clone(),
})
}
pub fn peer_info(&self) -> Option<&R::PeerInfo> {
self.info.get()
}
pub fn set_peer_info(&self, info: R::PeerInfo) {
if self.info.initialized() {
tracing::warn!("trying to set peer info, which is already initialized");
} else {
let _ = self.info.set(info);
}
}
pub fn is_transport_closed(&self) -> bool {
self.tx.is_closed()
}
}
#[derive(Debug)]
pub struct RunningService<R: ServiceRole, S: Service<R>> {
service: Arc<S>,
peer: Peer<R>,
handle: Option<tokio::task::JoinHandle<QuitReason>>,
cancellation_token: CancellationToken,
dg: DropGuard,
}
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()
}
#[inline]
pub fn cancellation_token(&self) -> RunningServiceCancellationToken {
RunningServiceCancellationToken(self.cancellation_token.clone())
}
/// Returns true if the service has been closed or cancelled.
#[inline]
pub fn is_closed(&self) -> bool {
self.handle.is_none() || self.cancellation_token.is_cancelled()
}
/// Wait for the service to complete.
///
/// This will block until the service loop terminates (either due to
/// cancellation, transport closure, or an error).
#[inline]
pub async fn waiting(mut self) -> Result<QuitReason, tokio::task::JoinError> {
match self.handle.take() {
Some(handle) => handle.await,
None => Ok(QuitReason::Closed),
}
}
/// Gracefully close the connection and wait for cleanup to complete.
///
/// This method cancels the service, waits for the background task to finish
/// (which includes calling `transport.close()`), and ensures all cleanup
/// operations complete before returning.
///
/// Unlike [`cancel`](Self::cancel), this method takes `&mut self` and can be
/// called without consuming the `RunningService`. After calling this method,
/// the service is considered closed and subsequent operations will fail.
///
/// # Example
///
/// ```rust,ignore
/// let mut client = ().serve(transport).await?;
/// // ... use the client ...
/// client.close().await?;
/// ```
pub async fn close(&mut self) -> Result<QuitReason, tokio::task::JoinError> {
if let Some(handle) = self.handle.take() {
// Disarm the drop guard so it doesn't try to cancel again
// We need to cancel manually and wait for completion
self.cancellation_token.cancel();
handle.await
} else {
// Already closed
Ok(QuitReason::Closed)
}
}
/// Gracefully close the connection with a timeout.
///
/// Similar to [`close`](Self::close), but returns after the specified timeout
/// if the cleanup doesn't complete in time. This is useful for ensuring
/// a bounded shutdown time.
///
/// Returns `Ok(Some(reason))` if shutdown completed within the timeout,
/// `Ok(None)` if the timeout was reached, or `Err` if there was a join error.
pub async fn close_with_timeout(
&mut self,
timeout: Duration,
) -> Result<Option<QuitReason>, tokio::task::JoinError> {
if let Some(handle) = self.handle.take() {
self.cancellation_token.cancel();
match tokio::time::timeout(timeout, handle).await {
Ok(result) => result.map(Some),
Err(_elapsed) => {
tracing::warn!(
"close_with_timeout: cleanup did not complete within {:?}",
timeout
);
Ok(None)
}
}
} else {
Ok(Some(QuitReason::Closed))
}
}
/// Cancel the service and wait for cleanup to complete.
///
/// This consumes the `RunningService` and ensures the connection is properly
/// closed. For a non-consuming alternative, see [`close`](Self::close).
pub async fn cancel(mut self) -> Result<QuitReason, tokio::task::JoinError> {
// Disarm the drop guard since we're handling cancellation explicitly
let _ = std::mem::replace(&mut self.dg, self.cancellation_token.clone().drop_guard());
self.close().await
}
}
impl<R: ServiceRole, S: Service<R>> Drop for RunningService<R, S> {
fn drop(&mut self) {
if self.handle.is_some() && !self.cancellation_token.is_cancelled() {
tracing::debug!(
"RunningService dropped without explicit close(). \
The connection will be closed asynchronously. \
For guaranteed cleanup, call close() or cancel() before dropping."
);
}
// The DropGuard will handle cancellation
}
}
// use a wrapper type so we can tweak the implementation if needed
pub struct RunningServiceCancellationToken(CancellationToken);
impl RunningServiceCancellationToken {
pub fn cancel(self) {
self.0.cancel();
}
}
#[derive(Debug)]
pub enum QuitReason {
Cancelled,
Closed,
JoinError(tokio::task::JoinError),
}
/// 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,
pub meta: Meta,
pub extensions: Extensions,
/// An interface to fetch the remote client or server
pub peer: Peer<R>,
}
/// Request execution context
#[derive(Debug, Clone)]
pub struct NotificationContext<R: ServiceRole> {
pub meta: Meta,
pub extensions: Extensions,
/// An interface to fetch the remote client or server
pub peer: Peer<R>,
}
/// Use this function to skip initialization process
pub fn serve_directly<R, S, T, E, A>(
service: S,
transport: T,
peer_info: Option<R::PeerInfo>,
) -> RunningService<R, S>
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())
}
/// Use this function to skip initialization process
pub fn serve_directly_with_ct<R, S, T, E, A>(
service: S,
transport: T,
peer_info: Option<R::PeerInfo>,
ct: CancellationToken,
) -> RunningService<R, S>
where
R: ServiceRole,
S: Service<R>,
T: IntoTransport<R, E, A>,
E: std::error::Error + Send + Sync + 'static,
{
let (peer, peer_rx) = Peer::new(Arc::new(AtomicU32RequestIdProvider::default()), peer_info);
serve_inner(service, transport.into_transport(), peer, peer_rx, ct)
}
#[instrument(skip_all)]
fn serve_inner<R, S, T>(
service: S,
transport: T,
peer: Peer<R>,
mut peer_rx: tokio::sync::mpsc::Receiver<PeerSinkMessage<R>>,
ct: CancellationToken,
) -> RunningService<R, S>
where
R: ServiceRole,
S: Service<R>,
T: Transport<R> + 'static,
{
const SINK_PROXY_BUFFER_SIZE: usize = 64;
let (sink_proxy_tx, mut sink_proxy_rx) =
tokio::sync::mpsc::channel::<TxJsonRpcMessage<R>>(SINK_PROXY_BUFFER_SIZE);
let peer_info = peer.peer_info();
if R::IS_CLIENT {
tracing::info!(?peer_info, "Service initialized as client");
} else {
tracing::info!(?peer_info, "Service initialized as server");
}
let mut local_responder_pool =
HashMap::<RequestId, Responder<Result<R::PeerResp, ServiceError>>>::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 current_span = tracing::Span::current();
let handle = tokio::spawn(async move {
let mut transport = transport.into_transport();
let mut batch_messages = VecDeque::<RxJsonRpcMessage<R>>::new();
let mut send_task_set = tokio::task::JoinSet::<SendTaskResult>::new();
#[derive(Debug)]
enum SendTaskResult {
Request {
id: RequestId,
result: Result<(), DynamicTransportError>,
},
Notification {
responder: Responder<Result<(), ServiceError>>,
cancellation_param: Option<CancelledNotificationParam>,
result: Result<(), DynamicTransportError>,
},
}
#[derive(Debug)]
enum Event<R: ServiceRole> {
ProxyMessage(PeerSinkMessage<R>),
PeerMessage(RxJsonRpcMessage<R>),
ToSink(TxJsonRpcMessage<R>),
SendTaskResult(SendTaskResult),
}
let quit_reason = loop {
let evt = if let Some(m) = batch_messages.pop_front() {
Event::PeerMessage(m)
} else {
tokio::select! {
m = sink_proxy_rx.recv(), if !sink_proxy_rx.is_closed() => {
if let Some(m) = m {
Event::ToSink(m)
} else {
continue
}
}
m = transport.receive() => {
if let Some(m) = m {
Event::PeerMessage(m)
} else {
// input stream closed
tracing::info!("input stream terminated");
break QuitReason::Closed
}
}
m = peer_rx.recv(), if !peer_rx.is_closed() => {
if let Some(m) = m {
Event::ProxyMessage(m)
} else {
continue
}
}
m = send_task_set.join_next(), if !send_task_set.is_empty() => {
let Some(result) = m else {
continue
};
match result {
Err(e) => {
// join error, which is serious, we should quit.
tracing::error!(%e, "send request task encounter a tokio join error");
break QuitReason::JoinError(e)
}
Ok(result) => {
Event::SendTaskResult(result)
}
}
}
_ = serve_loop_ct.cancelled() => {
tracing::info!("task cancelled");
break QuitReason::Cancelled
}
}
};
tracing::trace!(?evt, "new event");
match evt {
Event::SendTaskResult(SendTaskResult::Request { id, result }) => {
if let Err(e) = result {
if let Some(responder) = local_responder_pool.remove(&id) {
let _ = responder.send(Err(ServiceError::TransportSend(e)));
}
}
}
Event::SendTaskResult(SendTaskResult::Notification {
responder,
result,
cancellation_param,
}) => {
let response = if let Err(e) = result {
Err(ServiceError::TransportSend(e))
} else {
Ok(())
};
let _ = responder.send(response);
if let Some(param) = cancellation_param {
if let Some(responder) = local_responder_pool.remove(&param.request_id) {
tracing::info!(id = %param.request_id, reason = param.reason, "cancelled");
let _response_result = responder.send(Err(ServiceError::Cancelled {
reason: param.reason.clone(),
}));
}
}
}
// response and error
Event::ToSink(m) => {
if let Some(id) = match &m {
JsonRpcMessage::Response(response) => Some(&response.id),
JsonRpcMessage::Error(error) => Some(&error.id),
_ => None,
} {
if let Some(ct) = local_ct_pool.remove(id) {
ct.cancel();
}
let send = transport.send(m);
let current_span = tracing::Span::current();
tokio::spawn(async move {
let send_result = send.await;
if let Err(error) = send_result {
tracing::error!(%error, "fail to response message");
}
}.instrument(current_span));
}
}
Event::ProxyMessage(PeerSinkMessage::Request {
request,
id,
responder,
}) => {
local_responder_pool.insert(id.clone(), responder);
let send = transport.send(JsonRpcMessage::request(request, id.clone()));
{
let id = id.clone();
let current_span = tracing::Span::current();
send_task_set.spawn(send.map(move |r| SendTaskResult::Request {
id,
result: r.map_err(DynamicTransportError::new::<T, R>),
}).instrument(current_span));
}
}
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 = transport.send(JsonRpcMessage::notification(notification));
let current_span = tracing::Span::current();
send_task_set.spawn(send.map(move |result| SendTaskResult::Notification {
responder,
cancellation_param,
result: result.map_err(DynamicTransportError::new::<T, R>),
}).instrument(current_span));
}
Event::PeerMessage(JsonRpcMessage::Request(JsonRpcRequest {
id,
mut request,
..
})) => {
tracing::debug!(%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 mut extensions = Extensions::new();
let mut meta = Meta::new();
// avoid clone
// swap meta firstly, otherwise progress token will be lost
std::mem::swap(&mut meta, request.get_meta_mut());
std::mem::swap(&mut extensions, request.extensions_mut());
let context = RequestContext {
ct: context_ct,
id: id.clone(),
peer: peer.clone(),
meta,
extensions,
};
let current_span = tracing::Span::current();
tokio::spawn(async move {
let result = service
.handle_request(request, context)
.await;
let response = match result {
Ok(result) => {
tracing::debug!(%id, ?result, "response message");
JsonRpcMessage::response(result, id)
}
Err(error) => {
tracing::warn!(%id, ?error, "response error");
JsonRpcMessage::error(error, id)
}
};
let _send_result = sink.send(response).await;
}.instrument(current_span));
}
}
Event::PeerMessage(JsonRpcMessage::Notification(JsonRpcNotification {
notification,
..
})) => {
tracing::info!(?notification, "received notification");
// catch cancelled notification
let mut 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();
let mut extensions = Extensions::new();
let mut meta = Meta::new();
// avoid clone
std::mem::swap(&mut extensions, notification.extensions_mut());
std::mem::swap(&mut meta, notification.get_meta_mut());
let context = NotificationContext {
peer: peer.clone(),
meta,
extensions,
};
let current_span = tracing::Span::current();
tokio::spawn(async move {
let result = service.handle_notification(notification, context).await;
if let Err(error) = result {
tracing::warn!(%error, "Error sending notification");
}
}.instrument(current_span));
}
}
Event::PeerMessage(JsonRpcMessage::Response(JsonRpcResponse {
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(JsonRpcMessage::Error(JsonRpcError { 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 = transport.close().await;
if let Err(e) = sink_close_result {
tracing::error!(%e, "fail to close sink");
}
tracing::info!(?quit_reason, "serve finished");
quit_reason
}.instrument(current_span));
RunningService {
service,
peer: peer_return,
handle: Some(handle),
cancellation_token: ct.clone(),
dg: ct.drop_guard(),
}
}
codex-rs/vendor/rmcp/src/service/client.rs
Vendored
+543
View File
@@ -0,0 +1,543 @@
use std::borrow::Cow;
use thiserror::Error;
use super::*;
use crate::{
model::{
ArgumentInfo, CallToolRequest, CallToolRequestParams, CallToolResult,
CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage,
ClientNotification, ClientRequest, ClientResult, CompleteRequest, CompleteRequestParams,
CompleteResult, CompletionContext, CompletionInfo, ErrorData, GetPromptRequest,
GetPromptRequestParams, GetPromptResult, InitializeRequest, InitializedNotification,
JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest,
ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest,
ListToolsResult, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam,
ReadResourceRequest, ReadResourceRequestParams, ReadResourceResult, Reference, RequestId,
RootsListChangedNotification, ServerInfo, ServerJsonRpcMessage, ServerNotification,
ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, SubscribeRequest,
SubscribeRequestParams, UnsubscribeRequest, UnsubscribeRequestParams,
},
transport::DynamicTransportError,
};
/// It represents the error that may occur when serving the client.
///
/// if you want to handle the error, you can use `serve_client_with_ct` or `serve_client` with `Result<RunningService<RoleClient, S>, ClientError>`
#[derive(Error, Debug)]
pub enum ClientInitializeError {
#[error("expect initialized response, but received: {0:?}")]
ExpectedInitResponse(Option<ServerJsonRpcMessage>),
#[error("expect initialized result, but received: {0:?}")]
ExpectedInitResult(Option<ServerResult>),
#[error("conflict initialized response id: expected {0}, got {1}")]
ConflictInitResponseId(RequestId, RequestId),
#[error("connection closed: {0}")]
ConnectionClosed(String),
#[error("Send message error {error}, when {context}")]
TransportError {
error: DynamicTransportError,
context: Cow<'static, str>,
},
#[error("JSON-RPC error: {0}")]
JsonRpcError(ErrorData),
#[error("Cancelled")]
Cancelled,
}
impl ClientInitializeError {
pub fn transport<T: Transport<RoleClient> + 'static>(
error: T::Error,
context: impl Into<Cow<'static, str>>,
) -> Self {
Self::TransportError {
error: DynamicTransportError::new::<T, _>(error),
context: context.into(),
}
}
}
/// Helper function to get the next message from the stream
async fn expect_next_message<T>(
transport: &mut T,
context: &str,
) -> Result<ServerJsonRpcMessage, ClientInitializeError>
where
T: Transport<RoleClient>,
{
transport
.receive()
.await
.ok_or_else(|| ClientInitializeError::ConnectionClosed(context.to_string()))
}
/// Helper function to expect a response from the stream
async fn expect_response<T, S>(
transport: &mut T,
context: &str,
service: &S,
peer: Peer<RoleClient>,
) -> Result<(ServerResult, RequestId), ClientInitializeError>
where
T: Transport<RoleClient>,
S: Service<RoleClient>,
{
loop {
let message = expect_next_message(transport, context).await?;
match message {
// Expected message to complete the initialization
ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => {
break Ok((result, id));
}
// Handle JSON-RPC error responses
ServerJsonRpcMessage::Error(error) => {
break Err(ClientInitializeError::JsonRpcError(error.error));
}
// Server could send logging messages before handshake
ServerJsonRpcMessage::Notification(mut notification) => {
let ServerNotification::LoggingMessageNotification(logging) =
&mut notification.notification
else {
tracing::warn!(?notification, "Received unexpected message");
continue;
};
let mut context = NotificationContext {
peer: peer.clone(),
meta: Meta::default(),
extensions: Extensions::default(),
};
if let Some(meta) = logging.extensions.get_mut::<Meta>() {
std::mem::swap(&mut context.meta, meta);
}
std::mem::swap(&mut context.extensions, &mut logging.extensions);
if let Err(error) = service
.handle_notification(notification.notification, context)
.await
{
tracing::warn!(?error, "Handle logging before handshake failed.");
}
}
// Server could send pings before handshake
ServerJsonRpcMessage::Request(ref request)
if matches!(request.request, ServerRequest::PingRequest(_)) =>
{
tracing::trace!("Received ping request. Ignored.")
}
// Server SHOULD NOT send any other messages before handshake. We ignore them anyway
_ => tracing::warn!(?message, "Received unexpected message"),
}
}
}
#[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;
type InitializeError = ClientInitializeError;
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>, ClientInitializeError>> + Send
where
T: IntoTransport<RoleClient, E, A>,
E: std::error::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>, ClientInitializeError>
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
E: std::error::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>, ClientInitializeError>
where
S: Service<RoleClient>,
T: IntoTransport<RoleClient, E, A>,
E: std::error::Error + Send + Sync + 'static,
{
tokio::select! {
result = serve_client_with_ct_inner(service, transport.into_transport(), ct.clone()) => { result }
_ = ct.cancelled() => {
Err(ClientInitializeError::Cancelled)
}
}
}
async fn serve_client_with_ct_inner<S, T>(
service: S,
transport: T,
ct: CancellationToken,
) -> Result<RunningService<RoleClient, S>, ClientInitializeError>
where
S: Service<RoleClient>,
T: Transport<RoleClient> + 'static,
{
let mut transport = transport.into_transport();
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(),
extensions: Default::default(),
};
transport
.send(ClientJsonRpcMessage::request(
ClientRequest::InitializeRequest(init_request),
id.clone(),
))
.await
.map_err(|error| ClientInitializeError::TransportError {
error: DynamicTransportError::new::<T, _>(error),
context: "send initialize request".into(),
})?;
let (peer, peer_rx) = Peer::new(id_provider, None);
let (response, response_id) = expect_response(
&mut transport,
"initialize response",
&service,
peer.clone(),
)
.await?;
if id != response_id {
return Err(ClientInitializeError::ConflictInitResponseId(
id,
response_id,
));
}
let ServerResult::InitializeResult(initialize_result) = response else {
return Err(ClientInitializeError::ExpectedInitResult(Some(response)));
};
peer.set_peer_info(initialize_result);
// send notification
let notification = ClientJsonRpcMessage::notification(
ClientNotification::InitializedNotification(InitializedNotification {
method: Default::default(),
extensions: Default::default(),
}),
);
transport.send(notification).await.map_err(|error| {
ClientInitializeError::transport::<T>(error, "send initialized notification")
})?;
Ok(serve_inner(service, transport, peer, peer_rx, ct))
}
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,
extensions: 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: Option<$Param>) -> Result<$Resp, ServiceError> {
let result = self
.send_request(ClientRequest::$Req($Req {
method: Default::default(),
params,
extensions: Default::default(),
}))
.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,
extensions: Default::default(),
}))
.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,
extensions: Default::default(),
}))
.await?;
Ok(())
}
};
(peer_not $method:ident $Not:ident) => {
pub async fn $method(&self) -> Result<(), ServiceError> {
self.send_notification(ClientNotification::$Not($Not {
method: Default::default(),
extensions: Default::default(),
}))
.await?;
Ok(())
}
};
}
impl Peer<RoleClient> {
method!(peer_req complete CompleteRequest(CompleteRequestParams) => CompleteResult);
method!(peer_req set_level SetLevelRequest(SetLevelRequestParams));
method!(peer_req get_prompt GetPromptRequest(GetPromptRequestParams) => GetPromptResult);
method!(peer_req list_prompts ListPromptsRequest(PaginatedRequestParams)? => ListPromptsResult);
method!(peer_req list_resources ListResourcesRequest(PaginatedRequestParams)? => ListResourcesResult);
method!(peer_req list_resource_templates ListResourceTemplatesRequest(PaginatedRequestParams)? => ListResourceTemplatesResult);
method!(peer_req read_resource ReadResourceRequest(ReadResourceRequestParams) => ReadResourceResult);
method!(peer_req subscribe SubscribeRequest(SubscribeRequestParams) );
method!(peer_req unsubscribe UnsubscribeRequest(UnsubscribeRequestParams));
method!(peer_req call_tool CallToolRequest(CallToolRequestParams) => CallToolResult);
method!(peer_req list_tools ListToolsRequest(PaginatedRequestParams)? => 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(PaginatedRequestParams { meta: None, cursor }))
.await?;
tools.extend(result.tools);
cursor = result.next_cursor;
if cursor.is_none() {
break;
}
}
Ok(tools)
}
/// A wrapper method for [`Peer<RoleClient>::list_prompts`].
///
/// This function will call [`Peer<RoleClient>::list_prompts`] multiple times until all prompts 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(PaginatedRequestParams { meta: None, 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(PaginatedRequestParams { meta: None, 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(PaginatedRequestParams { meta: None, cursor }))
.await?;
resource_templates.extend(result.resource_templates);
cursor = result.next_cursor;
if cursor.is_none() {
break;
}
}
Ok(resource_templates)
}
/// Convenient method to get completion suggestions for a prompt argument
///
/// # Arguments
/// * `prompt_name` - Name of the prompt being completed
/// * `argument_name` - Name of the argument being completed
/// * `current_value` - Current partial value of the argument
/// * `context` - Optional context with previously resolved arguments
///
/// # Returns
/// CompletionInfo with suggestions for the specified prompt argument
pub async fn complete_prompt_argument(
&self,
prompt_name: impl Into<String>,
argument_name: impl Into<String>,
current_value: impl Into<String>,
context: Option<CompletionContext>,
) -> Result<CompletionInfo, ServiceError> {
let request = CompleteRequestParams {
meta: None,
r#ref: Reference::for_prompt(prompt_name),
argument: ArgumentInfo {
name: argument_name.into(),
value: current_value.into(),
},
context,
};
let result = self.complete(request).await?;
Ok(result.completion)
}
/// Convenient method to get completion suggestions for a resource URI argument
///
/// # Arguments
/// * `uri_template` - URI template pattern being completed
/// * `argument_name` - Name of the URI parameter being completed
/// * `current_value` - Current partial value of the parameter
/// * `context` - Optional context with previously resolved arguments
///
/// # Returns
/// CompletionInfo with suggestions for the specified resource URI argument
pub async fn complete_resource_argument(
&self,
uri_template: impl Into<String>,
argument_name: impl Into<String>,
current_value: impl Into<String>,
context: Option<CompletionContext>,
) -> Result<CompletionInfo, ServiceError> {
let request = CompleteRequestParams {
meta: None,
r#ref: Reference::for_resource(uri_template),
argument: ArgumentInfo {
name: argument_name.into(),
value: current_value.into(),
},
context,
};
let result = self.complete(request).await?;
Ok(result.completion)
}
/// Simple completion for a prompt argument without context
///
/// This is a convenience wrapper around `complete_prompt_argument` for
/// simple completion scenarios that don't require context awareness.
pub async fn complete_prompt_simple(
&self,
prompt_name: impl Into<String>,
argument_name: impl Into<String>,
current_value: impl Into<String>,
) -> Result<Vec<String>, ServiceError> {
let completion = self
.complete_prompt_argument(prompt_name, argument_name, current_value, None)
.await?;
Ok(completion.values)
}
/// Simple completion for a resource URI argument without context
///
/// This is a convenience wrapper around `complete_resource_argument` for
/// simple completion scenarios that don't require context awareness.
pub async fn complete_resource_simple(
&self,
uri_template: impl Into<String>,
argument_name: impl Into<String>,
current_value: impl Into<String>,
) -> Result<Vec<String>, ServiceError> {
let completion = self
.complete_resource_argument(uri_template, argument_name, current_value, None)
.await?;
Ok(completion.values)
}
}
codex-rs/vendor/rmcp/src/service/server.rs
Vendored
+895
View File
@@ -0,0 +1,895 @@
use std::borrow::Cow;
#[cfg(feature = "elicitation")]
use std::collections::HashSet;
use thiserror::Error;
#[cfg(feature = "elicitation")]
use url::Url;
use super::*;
#[cfg(feature = "elicitation")]
use crate::model::{
CreateElicitationRequest, CreateElicitationRequestParams, CreateElicitationResult,
ElicitationAction, ElicitationCompletionNotification, ElicitationResponseNotificationParam,
};
use crate::{
model::{
CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage,
ClientNotification, ClientRequest, ClientResult, CreateMessageRequest,
CreateMessageRequestParams, CreateMessageResult, ErrorData, ListRootsRequest,
ListRootsResult, LoggingMessageNotification, LoggingMessageNotificationParam,
ProgressNotification, ProgressNotificationParam, PromptListChangedNotification,
ProtocolVersion, ResourceListChangedNotification, ResourceUpdatedNotification,
ResourceUpdatedNotificationParam, ServerInfo, ServerNotification, ServerRequest,
ServerResult, ToolListChangedNotification,
},
transport::DynamicTransportError,
};
#[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;
type InitializeError = ServerInitializeError;
const IS_CLIENT: bool = false;
}
/// It represents the error that may occur when serving the server.
///
/// if you want to handle the error, you can use `serve_server_with_ct` or `serve_server` with `Result<RunningService<RoleServer, S>, ServerError>`
#[derive(Error, Debug)]
pub enum ServerInitializeError {
#[error("expect initialized request, but received: {0:?}")]
ExpectedInitializeRequest(Option<ClientJsonRpcMessage>),
#[error("expect initialized notification, but received: {0:?}")]
ExpectedInitializedNotification(Option<ClientJsonRpcMessage>),
#[error("connection closed: {0}")]
ConnectionClosed(String),
#[error("unexpected initialize result: {0:?}")]
UnexpectedInitializeResponse(ServerResult),
#[error("initialize failed: {0}")]
InitializeFailed(ErrorData),
#[error("unsupported protocol version: {0}")]
UnsupportedProtocolVersion(ProtocolVersion),
#[error("Send message error {error}, when {context}")]
TransportError {
error: DynamicTransportError,
context: Cow<'static, str>,
},
#[error("Cancelled")]
Cancelled,
}
impl ServerInitializeError {
pub fn transport<T: Transport<RoleServer> + 'static>(
error: T::Error,
context: impl Into<Cow<'static, str>>,
) -> Self {
Self::TransportError {
error: DynamicTransportError::new::<T, _>(error),
context: context.into(),
}
}
}
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>, ServerInitializeError>> + Send
where
T: IntoTransport<RoleServer, E, A>,
E: std::error::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>, ServerInitializeError>
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + Send + Sync + 'static,
{
serve_server_with_ct(service, transport, CancellationToken::new()).await
}
/// Helper function to get the next message from the stream
async fn expect_next_message<T>(
transport: &mut T,
context: &str,
) -> Result<ClientJsonRpcMessage, ServerInitializeError>
where
T: Transport<RoleServer>,
{
transport
.receive()
.await
.ok_or_else(|| ServerInitializeError::ConnectionClosed(context.to_string()))
}
/// Helper function to expect a request from the stream
async fn expect_request<T>(
transport: &mut T,
context: &str,
) -> Result<(ClientRequest, RequestId), ServerInitializeError>
where
T: Transport<RoleServer>,
{
let msg = expect_next_message(transport, context).await?;
let msg_clone = msg.clone();
msg.into_request()
.ok_or(ServerInitializeError::ExpectedInitializeRequest(Some(
msg_clone,
)))
}
/// Helper function to expect a notification from the stream
async fn expect_notification<T>(
transport: &mut T,
context: &str,
) -> Result<ClientNotification, ServerInitializeError>
where
T: Transport<RoleServer>,
{
let msg = expect_next_message(transport, context).await?;
let msg_clone = msg.clone();
msg.into_notification()
.ok_or(ServerInitializeError::ExpectedInitializedNotification(
Some(msg_clone),
))
}
pub async fn serve_server_with_ct<S, T, E, A>(
service: S,
transport: T,
ct: CancellationToken,
) -> Result<RunningService<RoleServer, S>, ServerInitializeError>
where
S: Service<RoleServer>,
T: IntoTransport<RoleServer, E, A>,
E: std::error::Error + Send + Sync + 'static,
{
tokio::select! {
result = serve_server_with_ct_inner(service, transport.into_transport(), ct.clone()) => { result }
_ = ct.cancelled() => {
Err(ServerInitializeError::Cancelled)
}
}
}
async fn serve_server_with_ct_inner<S, T>(
service: S,
transport: T,
ct: CancellationToken,
) -> Result<RunningService<RoleServer, S>, ServerInitializeError>
where
S: Service<RoleServer>,
T: Transport<RoleServer> + 'static,
{
let mut transport = transport.into_transport();
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
// Get initialize request
let (request, id) = expect_request(&mut transport, "initialized request").await?;
let ClientRequest::InitializeRequest(peer_info) = &request else {
return Err(ServerInitializeError::ExpectedInitializeRequest(Some(
ClientJsonRpcMessage::request(request, id),
)));
};
let (peer, peer_rx) = Peer::new(id_provider, Some(peer_info.params.clone()));
let context = RequestContext {
ct: ct.child_token(),
id: id.clone(),
meta: request.get_meta().clone(),
extensions: request.extensions().clone(),
peer: peer.clone(),
};
// Send initialize response
let init_response = service.handle_request(request.clone(), context).await;
let mut init_response = match init_response {
Ok(ServerResult::InitializeResult(init_response)) => init_response,
Ok(result) => {
return Err(ServerInitializeError::UnexpectedInitializeResponse(result));
}
Err(e) => {
transport
.send(ServerJsonRpcMessage::error(e.clone(), id))
.await
.map_err(|error| {
ServerInitializeError::transport::<T>(error, "sending error response")
})?;
return Err(ServerInitializeError::InitializeFailed(e));
}
};
let peer_protocol_version = peer_info.params.protocol_version.clone();
let protocol_version = match peer_protocol_version
.partial_cmp(&init_response.protocol_version)
.ok_or(ServerInitializeError::UnsupportedProtocolVersion(
peer_protocol_version,
))? {
std::cmp::Ordering::Less => peer_info.params.protocol_version.clone(),
_ => init_response.protocol_version,
};
init_response.protocol_version = protocol_version;
transport
.send(ServerJsonRpcMessage::response(
ServerResult::InitializeResult(init_response),
id,
))
.await
.map_err(|error| {
ServerInitializeError::transport::<T>(error, "sending initialize response")
})?;
// Wait for initialize notification
let notification = expect_notification(&mut transport, "initialize notification").await?;
let ClientNotification::InitializedNotification(_) = notification else {
return Err(ServerInitializeError::ExpectedInitializedNotification(
Some(ClientJsonRpcMessage::notification(notification)),
));
};
let context = NotificationContext {
meta: notification.get_meta().clone(),
extensions: notification.extensions().clone(),
peer: peer.clone(),
};
let _ = service.handle_notification(notification, context).await;
// Continue processing service
Ok(serve_inner(service, transport, peer, peer_rx, ct))
}
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(),
extensions: 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,
extensions: Default::default(),
}))
.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,
extensions: Default::default(),
}))
.await?;
Ok(())
}
};
(peer_not $method:ident $Not:ident) => {
pub async fn $method(&self) -> Result<(), ServiceError> {
self.send_notification(ServerNotification::$Not($Not {
method: Default::default(),
extensions: Default::default(),
}))
.await?;
Ok(())
}
};
// Timeout-only variants (base method should be created separately with peer_req)
(peer_req_with_timeout $method_with_timeout:ident $Req:ident() => $Resp: ident) => {
pub async fn $method_with_timeout(
&self,
timeout: Option<std::time::Duration>,
) -> Result<$Resp, ServiceError> {
let request = ServerRequest::$Req($Req {
method: Default::default(),
extensions: Default::default(),
});
let options = crate::service::PeerRequestOptions {
timeout,
meta: None,
};
let result = self
.send_request_with_option(request, options)
.await?
.await_response()
.await?;
match result {
ClientResult::$Resp(result) => Ok(result),
_ => Err(ServiceError::UnexpectedResponse),
}
}
};
(peer_req_with_timeout $method_with_timeout:ident $Req:ident($Param: ident) => $Resp: ident) => {
pub async fn $method_with_timeout(
&self,
params: $Param,
timeout: Option<std::time::Duration>,
) -> Result<$Resp, ServiceError> {
let request = ServerRequest::$Req($Req {
method: Default::default(),
params,
extensions: Default::default(),
});
let options = crate::service::PeerRequestOptions {
timeout,
meta: None,
};
let result = self
.send_request_with_option(request, options)
.await?
.await_response()
.await?;
match result {
ClientResult::$Resp(result) => Ok(result),
_ => Err(ServiceError::UnexpectedResponse),
}
}
};
}
impl Peer<RoleServer> {
/// Check if the client supports sampling tools capability.
pub fn supports_sampling_tools(&self) -> bool {
if let Some(client_info) = self.peer_info() {
client_info
.capabilities
.sampling
.as_ref()
.and_then(|s| s.tools.as_ref())
.is_some()
} else {
false
}
}
pub async fn create_message(
&self,
params: CreateMessageRequestParams,
) -> Result<CreateMessageResult, ServiceError> {
// MUST throw error when tools/toolChoice provided without capability
if (params.tools.is_some() || params.tool_choice.is_some())
&& !self.supports_sampling_tools()
{
return Err(ServiceError::McpError(ErrorData::invalid_params(
"tools or toolChoice provided but client does not support sampling tools capability",
None,
)));
}
// Validate message structure
params
.validate()
.map_err(|e| ServiceError::McpError(ErrorData::invalid_params(e, None)))?;
let result = self
.send_request(ServerRequest::CreateMessageRequest(CreateMessageRequest {
method: Default::default(),
params,
extensions: Default::default(),
}))
.await?;
match result {
ClientResult::CreateMessageResult(result) => Ok(*result),
_ => Err(ServiceError::UnexpectedResponse),
}
}
method!(peer_req list_roots ListRootsRequest() => ListRootsResult);
#[cfg(feature = "elicitation")]
method!(peer_req create_elicitation CreateElicitationRequest(CreateElicitationRequestParams) => CreateElicitationResult);
#[cfg(feature = "elicitation")]
method!(peer_req_with_timeout create_elicitation_with_timeout CreateElicitationRequest(CreateElicitationRequestParams) => CreateElicitationResult);
#[cfg(feature = "elicitation")]
method!(peer_not notify_url_elicitation_completed ElicitationCompletionNotification(ElicitationResponseNotificationParam));
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);
}
// =============================================================================
// ELICITATION CONVENIENCE METHODS
// These methods are specific to server role and provide typed elicitation functionality
// =============================================================================
/// Errors that can occur during typed elicitation operations
#[cfg(feature = "elicitation")]
#[derive(Error, Debug)]
pub enum ElicitationError {
/// The elicitation request failed at the service level
#[error("Service error: {0}")]
Service(#[from] ServiceError),
/// User explicitly declined to provide the requested information
/// This indicates a conscious decision by the user to reject the request
/// (e.g., clicked "Reject", "Decline", "No", etc.)
#[error("User explicitly declined the request")]
UserDeclined,
/// User dismissed the request without making an explicit choice
/// This indicates the user cancelled without explicitly declining
/// (e.g., closed dialog, clicked outside, pressed Escape, etc.)
#[error("User cancelled/dismissed the request")]
UserCancelled,
/// The response data could not be parsed into the requested type
#[error("Failed to parse response data: {error}\nReceived data: {data}")]
ParseError {
error: serde_json::Error,
data: serde_json::Value,
},
/// No response content was provided by the user
#[error("No response content provided")]
NoContent,
/// Client does not support elicitation capability
#[error("Client does not support elicitation - capability not declared during initialization")]
CapabilityNotSupported,
}
/// Marker trait to ensure that elicitation types generate object-type JSON schemas.
///
/// This trait provides compile-time safety to ensure that types used with
/// `elicit<T>()` methods will generate JSON schemas of type "object", which
/// aligns with MCP client expectations for structured data input.
///
/// # Type Safety Rationale
///
/// MCP clients typically expect JSON objects for elicitation schemas to
/// provide structured forms and validation. This trait prevents common
/// mistakes like:
///
/// ```compile_fail
/// // These would not compile due to missing ElicitationSafe bound:
/// let name: String = server.elicit("Enter name").await?; // Primitive
/// let items: Vec<i32> = server.elicit("Enter items").await?; // Array
/// ```
#[cfg(feature = "elicitation")]
pub trait ElicitationSafe: schemars::JsonSchema {}
/// Macro to mark types as safe for elicitation by verifying they generate object schemas.
///
/// This macro automatically implements the `ElicitationSafe` trait for struct types
/// that should be used with `elicit<T>()` methods.
///
/// # Example
///
/// ```rust
/// use rmcp::elicit_safe;
/// use schemars::JsonSchema;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Serialize, Deserialize, JsonSchema)]
/// struct UserProfile {
/// name: String,
/// email: String,
/// }
///
/// elicit_safe!(UserProfile);
///
/// // Now safe to use in async context:
/// // let profile: UserProfile = server.elicit("Enter profile").await?;
/// ```
#[cfg(feature = "elicitation")]
#[macro_export]
macro_rules! elicit_safe {
($($t:ty),* $(,)?) => {
$(
impl $crate::service::ElicitationSafe for $t {}
)*
};
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ElicitationMode {
Form,
Url,
}
#[cfg(feature = "elicitation")]
impl Peer<RoleServer> {
/// Check if the client supports elicitation capability
///
/// Returns true if the client declared elicitation capability during initialization,
/// false otherwise. According to MCP 2025-06-18 specification, clients that support
/// elicitation MUST declare the capability during initialization.
pub fn supported_elicitation_modes(&self) -> HashSet<ElicitationMode> {
if let Some(client_info) = self.peer_info() {
if let Some(elicit_capability) = &client_info.capabilities.elicitation {
let mut modes = HashSet::new();
// Backward compatibility: if neither form nor url is specified, assume form
if elicit_capability.form.is_none() && elicit_capability.url.is_none() {
modes.insert(ElicitationMode::Form);
} else {
if elicit_capability.form.is_some() {
modes.insert(ElicitationMode::Form);
}
if elicit_capability.url.is_some() {
modes.insert(ElicitationMode::Url);
}
}
modes
} else {
HashSet::new()
}
} else {
HashSet::new()
}
}
/// Request typed data from the user with automatic schema generation.
///
/// This method automatically generates the JSON schema from the Rust type using `schemars`,
/// eliminating the need to manually create schemas. The response is automatically parsed
/// into the requested type.
///
/// **Requires the `elicitation` feature to be enabled.**
///
/// # Type Requirements
/// The type `T` must implement:
/// - `schemars::JsonSchema` - for automatic schema generation
/// - `serde::Deserialize` - for parsing the response
///
/// # Arguments
/// * `message` - The prompt message for the user
///
/// # Returns
/// * `Ok(Some(data))` if user provided valid data that matches type T
/// * `Err(ElicitationError::UserDeclined)` if user explicitly declined the request
/// * `Err(ElicitationError::UserCancelled)` if user cancelled/dismissed the request
/// * `Err(ElicitationError::ParseError { .. })` if response data couldn't be parsed into type T
/// * `Err(ElicitationError::NoContent)` if no response content was provided
/// * `Err(ElicitationError::Service(_))` if the underlying service call failed
///
/// # Example
///
/// Add to your `Cargo.toml`:
/// ```toml
/// [dependencies]
/// rmcp = { version = "0.3", features = ["elicitation"] }
/// serde = { version = "1.0", features = ["derive"] }
/// schemars = "1.0"
/// ```
///
/// ```rust,no_run
/// # use rmcp::*;
/// # use rmcp::service::ElicitationError;
/// # use serde::{Deserialize, Serialize};
/// # use schemars::JsonSchema;
/// #
/// #[derive(Debug, Serialize, Deserialize, JsonSchema)]
/// struct UserProfile {
/// #[schemars(description = "Full name")]
/// name: String,
/// #[schemars(description = "Email address")]
/// email: String,
/// #[schemars(description = "Age")]
/// age: u8,
/// }
///
/// // Mark as safe for elicitation (generates object schema)
/// rmcp::elicit_safe!(UserProfile);
///
/// # async fn example(peer: Peer<RoleServer>) -> Result<(), Box<dyn std::error::Error>> {
/// match peer.elicit::<UserProfile>("Please enter your profile information").await {
/// Ok(Some(profile)) => {
/// println!("Name: {}, Email: {}, Age: {}", profile.name, profile.email, profile.age);
/// }
/// Ok(None) => {
/// println!("User provided no content");
/// }
/// Err(ElicitationError::UserDeclined) => {
/// println!("User explicitly declined to provide information");
/// // Handle explicit decline - perhaps offer alternatives
/// }
/// Err(ElicitationError::UserCancelled) => {
/// println!("User cancelled the request");
/// // Handle cancellation - perhaps prompt again later
/// }
/// Err(ElicitationError::ParseError { error, data }) => {
/// println!("Failed to parse response: {}\nData: {}", error, data);
/// }
/// Err(e) => return Err(e.into()),
/// }
/// # Ok(())
/// # }
/// ```
#[cfg(all(feature = "schemars", feature = "elicitation"))]
pub async fn elicit<T>(&self, message: impl Into<String>) -> Result<Option<T>, ElicitationError>
where
T: ElicitationSafe + for<'de> serde::Deserialize<'de>,
{
self.elicit_with_timeout(message, None).await
}
/// Request typed data from the user with custom timeout.
///
/// Same as `elicit()` but allows specifying a custom timeout for the request.
/// If the user doesn't respond within the timeout, the request will be cancelled.
///
/// # Arguments
/// * `message` - The prompt message for the user
/// * `timeout` - Optional timeout duration. If None, uses default timeout behavior
///
/// # Returns
/// Same as `elicit()` but may also return `ServiceError::Timeout` if timeout expires
///
/// # Example
/// ```rust,no_run
/// # use rmcp::*;
/// # use rmcp::service::ElicitationError;
/// # use serde::{Deserialize, Serialize};
/// # use schemars::JsonSchema;
/// # use std::time::Duration;
/// #
/// #[derive(Debug, Serialize, Deserialize, JsonSchema)]
/// struct QuickResponse {
/// answer: String,
/// }
///
/// // Mark as safe for elicitation
/// rmcp::elicit_safe!(QuickResponse);
///
/// # async fn example(peer: Peer<RoleServer>) -> Result<(), Box<dyn std::error::Error>> {
/// // Give user 30 seconds to respond
/// let timeout = Some(Duration::from_secs(30));
/// match peer.elicit_with_timeout::<QuickResponse>(
/// "Quick question - what's your answer?",
/// timeout
/// ).await {
/// Ok(Some(response)) => println!("Got answer: {}", response.answer),
/// Ok(None) => println!("User provided no content"),
/// Err(ElicitationError::UserDeclined) => {
/// println!("User explicitly declined");
/// // Handle explicit decline
/// }
/// Err(ElicitationError::UserCancelled) => {
/// println!("User cancelled/dismissed");
/// // Handle cancellation
/// }
/// Err(ElicitationError::Service(ServiceError::Timeout { .. })) => {
/// println!("User didn't respond in time");
/// }
/// Err(e) => return Err(e.into()),
/// }
/// # Ok(())
/// # }
/// ```
#[cfg(all(feature = "schemars", feature = "elicitation"))]
pub async fn elicit_with_timeout<T>(
&self,
message: impl Into<String>,
timeout: Option<std::time::Duration>,
) -> Result<Option<T>, ElicitationError>
where
T: ElicitationSafe + for<'de> serde::Deserialize<'de>,
{
// Check if client supports form elicitation capability
if !self
.supported_elicitation_modes()
.contains(&ElicitationMode::Form)
{
return Err(ElicitationError::CapabilityNotSupported);
}
// Generate schema automatically from type
let schema = crate::model::ElicitationSchema::from_type::<T>().map_err(|e| {
ElicitationError::Service(ServiceError::McpError(crate::ErrorData::invalid_params(
format!(
"Invalid schema for type {}: {}",
std::any::type_name::<T>(),
e
),
None,
)))
})?;
let response = self
.create_elicitation_with_timeout(
CreateElicitationRequestParams::FormElicitationParams {
meta: None,
message: message.into(),
requested_schema: schema,
},
timeout,
)
.await?;
match response.action {
crate::model::ElicitationAction::Accept => {
if let Some(value) = response.content {
match serde_json::from_value::<T>(value.clone()) {
Ok(parsed) => Ok(Some(parsed)),
Err(error) => Err(ElicitationError::ParseError { error, data: value }),
}
} else {
Err(ElicitationError::NoContent)
}
}
crate::model::ElicitationAction::Decline => Err(ElicitationError::UserDeclined),
crate::model::ElicitationAction::Cancel => Err(ElicitationError::UserCancelled),
}
}
/// Request the user to visit a URL and confirm completion.
///
/// This method sends a URL elicitation request to the client, prompting the user
/// to visit the specified URL and confirm completion. It returns the user's action
/// (accept/decline/cancel) without any additional data.
/// **Requires the `elicitation` feature to be enabled.**
///
/// # Arguments
/// * `message` - The prompt message for the user
/// * `url` - The URL the user is requested to visit
/// * `elicitation_id` - A unique identifier for this elicitation request
/// # Returns
/// * `Ok(action)` indicating the user's response action
/// * `Err(ElicitationError::CapabilityNotSupported)` if client does not support elicitation via URL
/// * `Err(ElicitationError::Service(_))` if the underlying service call failed
/// # Example
/// ```rust,no_run
/// # use rmcp::*;
/// # use rmcp::model::ElicitationAction;
/// # use url::Url;
///
/// async fn example(peer: Peer<RoleServer>) -> Result<(), Box<dyn std::error::Error>> {
/// let elicit_result = peer.elicit_url("Please visit the following URL to complete the action",
/// Url::parse("https://example.com/complete_action")?, "elicit_123").await?;
/// match elicit_result {
/// ElicitationAction::Accept => {
/// println!("User accepted and confirmed completion");
/// }
/// ElicitationAction::Decline => {
/// println!("User declined the request");
/// }
/// ElicitationAction::Cancel => {
/// println!("User cancelled/dismissed the request");
/// }
/// }
/// Ok(())
/// }
/// ```
#[cfg(feature = "elicitation")]
pub async fn elicit_url(
&self,
message: impl Into<String>,
url: impl Into<Url>,
elicitation_id: impl Into<String>,
) -> Result<ElicitationAction, ElicitationError> {
self.elicit_url_with_timeout(message, url, elicitation_id, None)
.await
}
/// Request the user to visit a URL and confirm completion.
///
/// Same as `elicit_url()` but allows specifying a custom timeout for the request.
///
/// # Arguments
/// * `message` - The prompt message for the user
/// * `url` - The URL the user is requested to visit
/// * `elicitation_id` - A unique identifier for this elicitation request
/// * `timeout` - Optional timeout duration. If None, uses default timeout behavior
/// # Returns
/// * `Ok(action)` indicating the user's response action
/// * `Err(ElicitationError::CapabilityNotSupported)` if client does not support elicitation via URL
/// * `Err(ElicitationError::Service(_))` if the underlying service call failed
/// # Example
/// ```rust,no_run
/// # use std::time::Duration;
/// use rmcp::*;
/// # use rmcp::model::ElicitationAction;
/// # use url::Url;
///
/// async fn example(peer: Peer<RoleServer>) -> Result<(), Box<dyn std::error::Error>> {
/// let elicit_result = peer.elicit_url_with_timeout("Please visit the following URL to complete the action",
/// Url::parse("https://example.com/complete_action")?,
/// "elicit_123",
/// Some(Duration::from_secs(30))).await?;
/// match elicit_result {
/// ElicitationAction::Accept => {
/// println!("User accepted and confirmed completion");
/// }
/// ElicitationAction::Decline => {
/// println!("User declined the request");
/// }
/// ElicitationAction::Cancel => {
/// println!("User cancelled/dismissed the request");
/// }
/// }
/// Ok(())
/// }
/// ```
#[cfg(feature = "elicitation")]
pub async fn elicit_url_with_timeout(
&self,
message: impl Into<String>,
url: impl Into<Url>,
elicitation_id: impl Into<String>,
timeout: Option<std::time::Duration>,
) -> Result<ElicitationAction, ElicitationError> {
// Check if client supports url elicitation
if !self
.supported_elicitation_modes()
.contains(&ElicitationMode::Url)
{
return Err(ElicitationError::CapabilityNotSupported);
}
let action = self
.create_elicitation_with_timeout(
CreateElicitationRequestParams::UrlElicitationParams {
meta: None,
message: message.into(),
url: url.into().to_string(),
elicitation_id: elicitation_id.into(),
},
timeout,
)
.await?
.action;
Ok(action)
}
}
codex-rs/vendor/rmcp/src/service/tower.rs
Vendored
+54
View File
@@ -0,0 +1,54 @@
use std::{future::poll_fn, marker::PhantomData};
use tower_service::Service as TowerService;
use super::NotificationContext;
use crate::service::{RequestContext, Service, ServiceRole};
pub struct TowerHandler<S, R: ServiceRole> {
pub service: S,
pub info: R::Info,
role: PhantomData<R>,
}
impl<S, R: ServiceRole> TowerHandler<S, R> {
pub fn new(service: S, info: R::Info) -> Self {
Self {
service,
role: PhantomData,
info,
}
}
}
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::ErrorData>,
S::Future: Send,
{
async fn handle_request(
&self,
request: R::PeerReq,
_context: RequestContext<R>,
) -> Result<R::Resp, crate::ErrorData> {
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,
_context: NotificationContext<R>,
) -> impl Future<Output = Result<(), crate::ErrorData>> + Send + '_ {
std::future::ready(Ok(()))
}
fn get_info(&self) -> R::Info {
self.info.clone()
}
}
codex-rs/vendor/rmcp/src/task_manager.rs
Vendored
+299
View File
@@ -0,0 +1,299 @@
use std::{any::Any, collections::HashMap, pin::Pin};
use futures::Future;
use tokio::{
sync::mpsc,
time::{Duration, timeout},
};
use crate::{
RoleServer,
error::{ErrorData as McpError, RmcpError as Error},
model::{CallToolResult, ClientRequest},
service::RequestContext,
};
/// Boxed future that represents an asynchronous operation managed by the processor.
pub type OperationFuture =
Pin<Box<dyn Future<Output = Result<Box<dyn OperationResultTransport>, Error>> + Send>>;
/// Describes metadata associated with an enqueued task.
#[derive(Debug, Clone)]
pub struct OperationDescriptor {
pub operation_id: String,
pub name: String,
pub client_request: Option<ClientRequest>,
pub context: Option<RequestContext<RoleServer>>,
pub ttl: Option<u64>,
}
impl OperationDescriptor {
pub fn new(operation_id: impl Into<String>, name: impl Into<String>) -> Self {
Self {
operation_id: operation_id.into(),
name: name.into(),
client_request: None,
context: None,
ttl: None,
}
}
pub fn with_client_request(mut self, request: ClientRequest) -> Self {
self.client_request = Some(request);
self
}
pub fn with_context(mut self, context: RequestContext<RoleServer>) -> Self {
self.context = Some(context);
self
}
pub fn with_ttl(mut self, ttl: u64) -> Self {
self.ttl = Some(ttl);
self
}
}
/// Operation message describing a unit of asynchronous work.
pub struct OperationMessage {
pub descriptor: OperationDescriptor,
pub future: OperationFuture,
}
impl OperationMessage {
pub fn new(descriptor: OperationDescriptor, future: OperationFuture) -> Self {
Self { descriptor, future }
}
}
/// Trait for operation result transport
pub trait OperationResultTransport: Send + Sync + 'static {
fn operation_id(&self) -> &String;
fn as_any(&self) -> &dyn std::any::Any;
}
// ===== Operation Processor =====
pub const DEFAULT_TASK_TIMEOUT_SECS: u64 = 300; // 5 minutes
/// Operation processor that coordinates extractors and handlers
pub struct OperationProcessor {
/// Currently running tasks keyed by id
running_tasks: HashMap<String, RunningTask>,
/// Completed results waiting to be collected
completed_results: Vec<TaskResult>,
task_result_receiver: mpsc::UnboundedReceiver<TaskResult>,
task_result_sender: mpsc::UnboundedSender<TaskResult>,
}
struct RunningTask {
task_handle: tokio::task::JoinHandle<()>,
started_at: std::time::Instant,
timeout: Option<u64>,
descriptor: OperationDescriptor,
}
pub struct TaskResult {
pub descriptor: OperationDescriptor,
pub result: Result<Box<dyn OperationResultTransport>, Error>,
}
/// Helper to generate an ISO 8601 timestamp for task metadata.
pub fn current_timestamp() -> String {
chrono::Utc::now().to_rfc3339()
}
/// Result transport for tool calls executed as tasks.
pub struct ToolCallTaskResult {
id: String,
pub result: Result<CallToolResult, McpError>,
}
impl ToolCallTaskResult {
pub fn new(id: impl Into<String>, result: Result<CallToolResult, McpError>) -> Self {
Self {
id: id.into(),
result,
}
}
}
impl OperationResultTransport for ToolCallTaskResult {
fn operation_id(&self) -> &String {
&self.id
}
fn as_any(&self) -> &dyn Any {
self
}
}
impl Default for OperationProcessor {
fn default() -> Self {
Self::new()
}
}
impl OperationProcessor {
pub fn new() -> Self {
let (task_result_sender, task_result_receiver) = mpsc::unbounded_channel();
Self {
running_tasks: HashMap::new(),
completed_results: Vec::new(),
task_result_receiver,
task_result_sender,
}
}
/// Submit an operation for asynchronous execution.
#[allow(clippy::result_large_err)]
pub fn submit_operation(&mut self, message: OperationMessage) -> Result<(), Error> {
if self
.running_tasks
.contains_key(&message.descriptor.operation_id)
{
return Err(Error::TaskError(format!(
"Operation with id {} is already running",
message.descriptor.operation_id
)));
}
self.spawn_async_task(message);
Ok(())
}
fn spawn_async_task(&mut self, message: OperationMessage) {
let OperationMessage { descriptor, future } = message;
let task_id = descriptor.operation_id.clone();
let timeout_secs = descriptor.ttl.or(Some(DEFAULT_TASK_TIMEOUT_SECS));
let sender = self.task_result_sender.clone();
let descriptor_for_result = descriptor.clone();
let timed_future = async move {
if let Some(secs) = timeout_secs {
match timeout(Duration::from_secs(secs), future).await {
Ok(result) => result,
Err(_) => Err(Error::TaskError("Operation timed out".to_string())),
}
} else {
future.await
}
};
let handle = tokio::spawn(async move {
let result = timed_future.await;
let task_result = TaskResult {
descriptor: descriptor_for_result,
result,
};
let _ = sender.send(task_result);
});
let running_task = RunningTask {
task_handle: handle,
started_at: std::time::Instant::now(),
timeout: timeout_secs,
descriptor,
};
self.running_tasks.insert(task_id, running_task);
}
/// Collect completed results from running tasks and remove them from the running tasks map.
fn collect_completed_results(&mut self) {
while let Ok(result) = self.task_result_receiver.try_recv() {
self.running_tasks.remove(&result.descriptor.operation_id);
self.completed_results.push(result);
}
}
/// Check for tasks that have exceeded their timeout and handle them appropriately.
pub fn check_timeouts(&mut self) {
self.collect_completed_results();
let now = std::time::Instant::now();
let mut timed_out_tasks = Vec::new();
for (task_id, task) in &self.running_tasks {
if let Some(timeout_duration) = task.timeout {
if now.duration_since(task.started_at).as_secs() > timeout_duration {
task.task_handle.abort();
timed_out_tasks.push(task_id.clone());
}
}
}
for task_id in timed_out_tasks {
if let Some(task) = self.running_tasks.remove(&task_id) {
let timeout_result = TaskResult {
descriptor: task.descriptor,
result: Err(Error::TaskError("Operation timed out".to_string())),
};
self.completed_results.push(timeout_result);
}
}
}
/// Get the number of running tasks.
pub fn running_task_count(&mut self) -> usize {
self.collect_completed_results();
self.running_tasks.len()
}
/// Cancel all running tasks.
pub fn cancel_all_tasks(&mut self) {
for (_, task) in self.running_tasks.drain() {
task.task_handle.abort();
}
while self.task_result_receiver.try_recv().is_ok() {}
self.completed_results.clear();
}
/// List running task ids.
pub fn list_running(&mut self) -> Vec<String> {
self.collect_completed_results();
self.running_tasks.keys().cloned().collect()
}
/// Returns a snapshot of completed task results.
pub fn peek_completed(&mut self) -> &[TaskResult] {
self.collect_completed_results();
&self.completed_results
}
/// Fetch the metadata for a running or recently completed task.
pub fn task_descriptor(&self, task_id: &str) -> Option<&OperationDescriptor> {
if let Some(task) = self.running_tasks.get(task_id) {
return Some(&task.descriptor);
}
self.completed_results
.iter()
.rev()
.find(|result| result.descriptor.operation_id == task_id)
.map(|result| &result.descriptor)
}
/// Attempt to cancel a running task.
pub fn cancel_task(&mut self, task_id: &str) -> bool {
self.collect_completed_results();
if let Some(task) = self.running_tasks.remove(task_id) {
task.task_handle.abort();
// Insert a cancelled result so callers can observe the terminal state.
let cancel_result = TaskResult {
descriptor: task.descriptor,
result: Err(Error::TaskError("Operation cancelled".to_string())),
};
self.completed_results.push(cancel_result);
return true;
}
false
}
/// Retrieve a completed task result if available.
pub fn take_completed_result(&mut self, task_id: &str) -> Option<TaskResult> {
self.collect_completed_results();
if let Some(position) = self
.completed_results
.iter()
.position(|result| result.descriptor.operation_id == task_id)
{
Some(self.completed_results.remove(position))
} else {
None
}
}
}
codex-rs/vendor/rmcp/src/transport.rs
Vendored
+277
View File
@@ -0,0 +1,277 @@
//! # Transport
//! The transport type must implemented [`Transport`] trait, which allow it send message concurrently and receive message sequentially.
//！
//! ## Standard Transport Types
//! There are 2 pairs of standard transport types:
//!
//! | transport | client | server |
//! |:-: |:-: |:-: |
//! | std IO | [`child_process::TokioChildProcess`] | [`io::stdio`] |
//! | streamable http | [`streamable_http_client::StreamableHttpClientTransport`] | [`streamable_http_server::StreamableHttpService`] |
//!
//！## Helper Transport Types
//! Thers are several helper transport types that can help you to create transport quickly.
//!
//! ### [Worker Transport](`worker::WorkerTransport`)
//! Which allows you to run a worker and process messages in another tokio task.
//!
//! ### [Async Read/Write Transport](`async_rw::AsyncRwTransport`)
//! You need to enable `transport-async-rw` feature to use this transport.
//!
//! This transport is used to create a transport from a byte stream which implemented [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`].
//!
//! This could be very helpful when you want to create a transport from a byte stream, such as a file or a tcp connection.
//!
//! ### [Sink/Stream Transport](`sink_stream::SinkStreamTransport`)
//! This transport is used to create a transport from a sink and a stream.
//!
//! This could be very helpful when you want to create a transport from a duplex object stream, such as a websocket connection.
//!
//! ## [IntoTransport](`IntoTransport`) trait
//! [`IntoTransport`] is a helper trait that implicitly convert a type into a transport type.
//!
//! ### These types is automatically implemented [`IntoTransport`] trait
//! 1. A type that already implement both [`futures::Sink`] and [`futures::Stream`] trait, or a tuple `(Tx, Rx)` where `Tx` is [`futures::Sink`] and `Rx` is [`futures::Stream`].
//! 2. A type that implement both [`tokio::io::AsyncRead`] and [`tokio::io::AsyncWrite`] trait. or a tuple `(R, W)` where `R` is [`tokio::io::AsyncRead`] and `W` is [`tokio::io::AsyncWrite`].
//! 3. A type that implement [Worker](`worker::Worker`) trait.
//! 4. A type that implement [`Transport`] trait.
//!
//! ## Examples
//!
//! ```rust
//! # use rmcp::{
//! # ServiceExt, serve_server,
//! # };
//! #[cfg(feature = "client")]
//! #[cfg_attr(docsrs, doc(cfg(feature = "client")))]
//! # use rmcp::serve_client;
//!
//! // create transport from tcp stream
//! #[cfg(feature = "client")]
//! #[cfg_attr(docsrs, doc(cfg(feature = "client")))]
//! async fn client() -> Result<(), Box<dyn std::error::Error>> {
//! let stream = tokio::net::TcpSocket::new_v4()?
//! .connect("127.0.0.1:8001".parse()?)
//! .await?;
//! let client = ().serve(stream).await?;
//! let tools = client.peer().list_tools(Default::default()).await?;
//! println!("{:?}", tools);
//! Ok(())
//! }
//!
//! // create transport from std io
//! #[cfg(feature = "client")]
//! #[cfg_attr(docsrs, doc(cfg(feature = "client")))]
//! async fn io() -> Result<(), Box<dyn std::error::Error>> {
//! let client = ().serve((tokio::io::stdin(), tokio::io::stdout())).await?;
//! let tools = client.peer().list_tools(Default::default()).await?;
//! println!("{:?}", tools);
//! Ok(())
//! }
//! ```
use std::{borrow::Cow, sync::Arc};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
pub mod sink_stream;
#[cfg(feature = "transport-async-rw")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-async-rw")))]
pub mod async_rw;
#[cfg(feature = "transport-worker")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-worker")))]
pub mod worker;
#[cfg(feature = "transport-worker")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-worker")))]
pub use worker::WorkerTransport;
#[cfg(feature = "transport-child-process")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-child-process")))]
pub mod child_process;
#[cfg(feature = "transport-child-process")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-child-process")))]
pub use child_process::{ConfigureCommandExt, TokioChildProcess};
#[cfg(feature = "transport-io")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-io")))]
pub mod io;
#[cfg(feature = "transport-io")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-io")))]
pub use io::stdio;
#[cfg(feature = "auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
pub mod auth;
#[cfg(feature = "auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
pub use auth::{
AuthError, AuthorizationManager, AuthorizationSession, AuthorizedHttpClient, CredentialStore,
InMemoryCredentialStore, InMemoryStateStore, StateStore, StoredAuthorizationState,
StoredCredentials,
};
// #[cfg(feature = "transport-ws")]
// #[cfg_attr(docsrs, doc(cfg(feature = "transport-ws")))]
// pub mod ws;
#[cfg(feature = "transport-streamable-http-server-session")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server-session")))]
pub mod streamable_http_server;
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub use streamable_http_server::tower::{StreamableHttpServerConfig, StreamableHttpService};
#[cfg(feature = "transport-streamable-http-client")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client")))]
pub mod streamable_http_client;
#[cfg(feature = "transport-streamable-http-client")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client")))]
pub use streamable_http_client::StreamableHttpClientTransport;
/// Common use codes
pub mod common;
pub trait Transport<R>: Send
where
R: ServiceRole,
{
type Error: std::error::Error + Send + Sync + 'static;
fn name() -> Cow<'static, str> {
std::any::type_name::<Self>().into()
}
/// Send a message to the transport
///
/// Notice that the future returned by this function should be `Send` and `'static`.
/// It's because the sending message could be executed concurrently.
///
fn send(
&mut self,
item: TxJsonRpcMessage<R>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static;
/// Receive a message from the transport, this operation is sequential.
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<R>>> + Send;
/// Close the transport
fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send;
}
pub trait IntoTransport<R, E, A>: Send + 'static
where
R: ServiceRole,
E: std::error::Error + Send + 'static,
{
fn into_transport(self) -> impl Transport<R, Error = E> + 'static;
}
pub enum TransportAdapterIdentity {}
impl<R, T, E> IntoTransport<R, E, TransportAdapterIdentity> for T
where
T: Transport<R, Error = E> + Send + 'static,
R: ServiceRole,
E: std::error::Error + Send + Sync + 'static,
{
fn into_transport(self) -> impl Transport<R, Error = E> + 'static {
self
}
}
/// A transport that can send a single message and then close itself
pub struct OneshotTransport<R>
where
R: ServiceRole,
{
message: Option<RxJsonRpcMessage<R>>,
sender: tokio::sync::mpsc::Sender<TxJsonRpcMessage<R>>,
termination: Arc<tokio::sync::Semaphore>,
}
impl<R> OneshotTransport<R>
where
R: ServiceRole,
{
pub fn new(
message: RxJsonRpcMessage<R>,
) -> (Self, tokio::sync::mpsc::Receiver<TxJsonRpcMessage<R>>) {
let (sender, receiver) = tokio::sync::mpsc::channel(16);
(
Self {
message: Some(message),
sender,
termination: Arc::new(tokio::sync::Semaphore::new(0)),
},
receiver,
)
}
}
impl<R> Transport<R> for OneshotTransport<R>
where
R: ServiceRole,
{
type Error = tokio::sync::mpsc::error::SendError<TxJsonRpcMessage<R>>;
fn send(
&mut self,
item: TxJsonRpcMessage<R>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
let sender = self.sender.clone();
let terminate = matches!(item, TxJsonRpcMessage::<R>::Response(_))
|| matches!(item, TxJsonRpcMessage::<R>::Error(_));
let termination = self.termination.clone();
async move {
sender.send(item).await?;
if terminate {
termination.add_permits(1);
}
Ok(())
}
}
async fn receive(&mut self) -> Option<RxJsonRpcMessage<R>> {
if let Some(msg) = self.message.take() {
return Some(msg);
}
let _ = self.termination.acquire().await;
None
}
fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send {
self.message.take();
std::future::ready(Ok(()))
}
}
#[derive(Debug, thiserror::Error)]
#[error("Transport [{transport_name}] error: {error}")]
pub struct DynamicTransportError {
pub transport_name: Cow<'static, str>,
pub transport_type_id: std::any::TypeId,
#[source]
pub error: Box<dyn std::error::Error + Send + Sync>,
}
impl DynamicTransportError {
pub fn new<T: Transport<R> + 'static, R: ServiceRole>(e: T::Error) -> Self {
Self {
transport_name: T::name(),
transport_type_id: std::any::TypeId::of::<T>(),
error: Box::new(e),
}
}
pub fn downcast<T: Transport<R> + 'static, R: ServiceRole>(self) -> Result<T::Error, Self> {
if !self.is::<T, R>() {
Err(self)
} else {
Ok(self
.error
.downcast::<T::Error>()
.map(|e| *e)
.expect("type is checked"))
}
}
pub fn is<T: Transport<R> + 'static, R: ServiceRole>(&self) -> bool {
self.error.is::<T::Error>() && self.transport_type_id == std::any::TypeId::of::<T>()
}
}
codex-rs/vendor/rmcp/src/transport/async_rw.rs
Vendored
+557
View File
@@ -0,0 +1,557 @@
use std::{marker::PhantomData, sync::Arc};
// use crate::schema::*;
use futures::{SinkExt, StreamExt};
use serde::{Serialize, de::DeserializeOwned};
use thiserror::Error;
use tokio::{
io::{AsyncRead, AsyncWrite},
sync::Mutex,
};
use tokio_util::{
bytes::{Buf, BufMut, BytesMut},
codec::{Decoder, Encoder, FramedRead, FramedWrite},
};
use super::{IntoTransport, Transport};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
pub enum TransportAdapterAsyncRW {}
impl<Role, R, W> IntoTransport<Role, std::io::Error, TransportAdapterAsyncRW> for (R, W)
where
Role: ServiceRole,
R: AsyncRead + Send + 'static + Unpin,
W: AsyncWrite + Send + 'static + Unpin,
{
fn into_transport(self) -> impl Transport<Role, Error = std::io::Error> + 'static {
AsyncRwTransport::new(self.0, self.1)
}
}
pub enum TransportAdapterAsyncCombinedRW {}
impl<Role, S> IntoTransport<Role, std::io::Error, TransportAdapterAsyncCombinedRW> for S
where
Role: ServiceRole,
S: AsyncRead + AsyncWrite + Send + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = std::io::Error> + 'static {
IntoTransport::<Role, std::io::Error, TransportAdapterAsyncRW>::into_transport(
tokio::io::split(self),
)
}
}
pub type TransportWriter<Role, W> = FramedWrite<W, JsonRpcMessageCodec<TxJsonRpcMessage<Role>>>;
pub struct AsyncRwTransport<Role: ServiceRole, R: AsyncRead, W: AsyncWrite> {
read: FramedRead<R, JsonRpcMessageCodec<RxJsonRpcMessage<Role>>>,
write: Arc<Mutex<Option<TransportWriter<Role, W>>>>,
}
impl<Role: ServiceRole, R, W> AsyncRwTransport<Role, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
pub fn new(read: R, write: W) -> Self {
let read = FramedRead::new(
read,
JsonRpcMessageCodec::<RxJsonRpcMessage<Role>>::default(),
);
let write = Arc::new(Mutex::new(Some(FramedWrite::new(
write,
JsonRpcMessageCodec::<TxJsonRpcMessage<Role>>::default(),
))));
Self { read, write }
}
}
#[cfg(feature = "client")]
#[cfg_attr(docsrs, doc(cfg(feature = "client")))]
impl<R, W> AsyncRwTransport<crate::RoleClient, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
pub fn new_client(read: R, write: W) -> Self {
Self::new(read, write)
}
}
#[cfg(feature = "server")]
#[cfg_attr(docsrs, doc(cfg(feature = "server")))]
impl<R, W> AsyncRwTransport<crate::RoleServer, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
pub fn new_server(read: R, write: W) -> Self {
Self::new(read, write)
}
}
impl<Role: ServiceRole, R, W> Transport<Role> for AsyncRwTransport<Role, R, W>
where
R: Send + AsyncRead + Unpin,
W: Send + AsyncWrite + Unpin + 'static,
{
type Error = std::io::Error;
fn send(
&mut self,
item: TxJsonRpcMessage<Role>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
let lock = self.write.clone();
async move {
let mut write = lock.lock().await;
if let Some(ref mut write) = *write {
write.send(item).await.map_err(Into::into)
} else {
Err(std::io::Error::new(
std::io::ErrorKind::NotConnected,
"Transport is closed",
))
}
}
}
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<Role>>> {
let next = self.read.next();
async {
next.await.and_then(|e| {
e.inspect_err(|e| {
tracing::error!("Error reading from stream: {}", e);
})
.ok()
})
}
}
async fn close(&mut self) -> Result<(), Self::Error> {
let mut write = self.write.lock().await;
drop(write.take());
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct JsonRpcMessageCodec<T> {
_marker: PhantomData<fn() -> T>,
next_index: usize,
max_length: usize,
is_discarding: bool,
}
impl<T> Default for JsonRpcMessageCodec<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> JsonRpcMessageCodec<T> {
pub fn new() -> Self {
Self {
_marker: PhantomData,
next_index: 0,
max_length: usize::MAX,
is_discarding: false,
}
}
pub fn new_with_max_length(max_length: usize) -> Self {
Self {
max_length,
..Self::new()
}
}
pub fn max_length(&self) -> usize {
self.max_length
}
}
fn without_carriage_return(s: &[u8]) -> &[u8] {
if let Some(&b'\r') = s.last() {
&s[..s.len() - 1]
} else {
s
}
}
/// Check if a method is a standard MCP method (request, response, or notification).
/// This includes both requests and notifications defined in the MCP specification.
///
/// Based on MCP specification 2025-06-18: https://modelcontextprotocol.io/specification/2025-06-18
fn is_standard_method(method: &str) -> bool {
matches!(
method,
"initialize"
| "ping"
| "prompts/get"
| "prompts/list"
| "resources/list"
| "resources/read"
| "resources/subscribe"
| "resources/unsubscribe"
| "resources/templates/list"
| "tools/call"
| "tools/list"
| "completion/complete"
| "logging/setLevel"
| "roots/list"
| "sampling/createMessage"
) || is_standard_notification(method)
}
fn is_standard_notification(method: &str) -> bool {
matches!(
method,
"notifications/cancelled"
| "notifications/initialized"
| "notifications/message"
| "notifications/progress"
| "notifications/prompts/list_changed"
| "notifications/resources/list_changed"
| "notifications/resources/updated"
| "notifications/roots/list_changed"
| "notifications/tools/list_changed"
)
}
/// Determines if a notification should be ignored for compatibility.
fn should_ignore_notification(json_value: &serde_json::Value, method: &str) -> bool {
let is_notification = json_value.get("id").is_none();
// Ignore non-MCP notifications (like LSP messages) for compatibility
if is_notification && !is_standard_method(method) {
tracing::trace!(
"Ignoring non-MCP notification '{}' for compatibility",
method
);
return true;
}
// Ignore non-standard MCP notifications
matches!(
(
method.starts_with("notifications/"),
is_standard_notification(method)
),
(true, false)
)
}
/// Try to parse a message with compatibility handling for non-standard notifications
fn try_parse_with_compatibility<T: serde::de::DeserializeOwned>(
line: &[u8],
context: &str,
) -> Result<Option<T>, JsonRpcMessageCodecError> {
if let Ok(line_str) = std::str::from_utf8(line) {
match serde_json::from_slice(line) {
Ok(item) => Ok(Some(item)),
Err(e) => {
// Check if this is a notification that should be ignored for compatibility
if let Ok(json_value) = serde_json::from_str::<serde_json::Value>(line_str) {
if let Some(method) =
json_value.get("method").and_then(serde_json::Value::as_str)
{
if should_ignore_notification(&json_value, method) {
return Ok(None);
}
}
}
tracing::debug!(
"Failed to parse message {}: {} | Error: {}",
context,
line_str,
e
);
Err(JsonRpcMessageCodecError::Serde(e))
}
}
} else {
serde_json::from_slice(line)
.map(Some)
.map_err(JsonRpcMessageCodecError::Serde)
}
}
#[derive(Debug, Error)]
pub enum JsonRpcMessageCodecError {
#[error("max line length exceeded")]
MaxLineLengthExceeded,
#[error("serde error {0}")]
Serde(#[from] serde_json::Error),
#[error("io error {0}")]
Io(#[from] std::io::Error),
}
impl From<JsonRpcMessageCodecError> for std::io::Error {
fn from(value: JsonRpcMessageCodecError) -> Self {
match value {
JsonRpcMessageCodecError::MaxLineLengthExceeded => {
std::io::Error::new(std::io::ErrorKind::InvalidData, value)
}
JsonRpcMessageCodecError::Serde(e) => e.into(),
JsonRpcMessageCodecError::Io(e) => e,
}
}
}
impl<T: DeserializeOwned> Decoder for JsonRpcMessageCodec<T> {
type Item = T;
type Error = JsonRpcMessageCodecError;
fn decode(
&mut self,
buf: &mut BytesMut,
) -> Result<Option<Self::Item>, JsonRpcMessageCodecError> {
loop {
// Determine how far into the buffer we'll search for a newline. If
// there's no max_length set, we'll read to the end of the buffer.
let read_to = std::cmp::min(self.max_length.saturating_add(1), buf.len());
let newline_offset = buf[self.next_index..read_to]
.iter()
.position(|b| *b == b'\n');
match (self.is_discarding, newline_offset) {
(true, Some(offset)) => {
// If we found a newline, discard up to that offset and
// then stop discarding. On the next iteration, we'll try
// to read a line normally.
buf.advance(offset + self.next_index + 1);
self.is_discarding = false;
self.next_index = 0;
}
(true, None) => {
// Otherwise, we didn't find a newline, so we'll discard
// everything we read. On the next iteration, we'll continue
// discarding up to max_len bytes unless we find a newline.
buf.advance(read_to);
self.next_index = 0;
if buf.is_empty() {
return Ok(None);
}
}
(false, Some(offset)) => {
// Found a line!
let newline_index = offset + self.next_index;
self.next_index = 0;
let line = buf.split_to(newline_index + 1);
let line = &line[..line.len() - 1];
let line = without_carriage_return(line);
// Use compatibility handling function
let item = match try_parse_with_compatibility(line, "decode")? {
Some(item) => item,
None => return Ok(None), // Skip non-standard message
};
return Ok(Some(item));
}
(false, None) if buf.len() > self.max_length => {
// Reached the maximum length without finding a
// newline, return an error and start discarding on the
// next call.
self.is_discarding = true;
return Err(JsonRpcMessageCodecError::MaxLineLengthExceeded);
}
(false, None) => {
// We didn't find a line or reach the length limit, so the next
// call will resume searching at the current offset.
self.next_index = read_to;
return Ok(None);
}
}
}
}
fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<T>, JsonRpcMessageCodecError> {
Ok(match self.decode(buf)? {
Some(frame) => Some(frame),
None => {
self.next_index = 0;
// No terminating newline - return remaining data, if any
if buf.is_empty() || buf == &b"\r"[..] {
None
} else {
let line = buf.split_to(buf.len());
let line = without_carriage_return(&line);
// Use compatibility handling function
let item = match try_parse_with_compatibility(line, "decode_eof")? {
Some(item) => item,
None => return Ok(None), // Skip non-standard message
};
Some(item)
}
}
})
}
}
impl<T: Serialize> Encoder<T> for JsonRpcMessageCodec<T> {
type Error = JsonRpcMessageCodecError;
fn encode(&mut self, item: T, buf: &mut BytesMut) -> Result<(), JsonRpcMessageCodecError> {
serde_json::to_writer(buf.writer(), &item)?;
buf.put_u8(b'\n');
Ok(())
}
}
#[cfg(test)]
mod test {
use futures::{Sink, Stream};
use super::*;
fn from_async_read<T: DeserializeOwned, R: AsyncRead>(reader: R) -> impl Stream<Item = T> {
FramedRead::new(reader, JsonRpcMessageCodec::<T>::default()).filter_map(|result| {
if let Err(e) = &result {
tracing::error!("Error reading from stream: {}", e);
}
futures::future::ready(result.ok())
})
}
fn from_async_write<T: Serialize, W: AsyncWrite + Send>(
writer: W,
) -> impl Sink<T, Error = std::io::Error> {
FramedWrite::new(writer, JsonRpcMessageCodec::<T>::default()).sink_map_err(Into::into)
}
#[tokio::test]
async fn test_decode() {
use futures::StreamExt;
use tokio::io::BufReader;
let data = r#"{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":1}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":2}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":3}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":4}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":5}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":6}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":7}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":8}
{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":9}
{"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":10}
"#;
let mut cursor = BufReader::new(data.as_bytes());
let mut stream = from_async_read::<serde_json::Value, _>(&mut cursor);
for i in 1..=10 {
let item = stream.next().await.unwrap();
assert_eq!(
item,
serde_json::json!({
"jsonrpc": "2.0",
"method": "subtract",
"params": if i % 2 != 0 { [42, 23] } else { [23, 42] },
"id": i,
})
);
}
}
#[tokio::test]
async fn test_encode() {
let test_messages = vec![
serde_json::json!({
"jsonrpc": "2.0",
"method": "subtract",
"params": [42, 23],
"id": 1,
}),
serde_json::json!({
"jsonrpc": "2.0",
"method": "subtract",
"params": [23, 42],
"id": 2,
}),
];
// Create a buffer to write to
let mut buffer = Vec::new();
let mut writer = from_async_write(&mut buffer);
// Write the test messages
for message in test_messages.iter() {
writer.send(message.clone()).await.unwrap();
}
writer.close().await.unwrap();
drop(writer);
// Parse the buffer back into lines and check each one
let output = String::from_utf8_lossy(&buffer);
let mut lines = output.lines();
for expected_message in test_messages {
let line = lines.next().unwrap();
let parsed_message: serde_json::Value = serde_json::from_str(line).unwrap();
assert_eq!(parsed_message, expected_message);
}
// Make sure there are no extra lines
assert!(lines.next().is_none());
}
#[test]
fn test_standard_notification_check() {
// Test that all standard notifications are recognized
assert!(is_standard_notification("notifications/cancelled"));
assert!(is_standard_notification("notifications/initialized"));
assert!(is_standard_notification("notifications/progress"));
assert!(is_standard_notification(
"notifications/resources/list_changed"
));
assert!(is_standard_notification("notifications/resources/updated"));
assert!(is_standard_notification(
"notifications/prompts/list_changed"
));
assert!(is_standard_notification("notifications/tools/list_changed"));
assert!(is_standard_notification("notifications/message"));
assert!(is_standard_notification("notifications/roots/list_changed"));
// Test that non-standard notifications are not recognized
assert!(!is_standard_notification("notifications/stderr"));
assert!(!is_standard_notification("notifications/custom"));
assert!(!is_standard_notification("notifications/debug"));
assert!(!is_standard_notification("some/other/method"));
}
#[test]
fn test_compatibility_function() {
// Test the compatibility function directly
let stderr_message =
r#"{"method":"notifications/stderr","params":{"content":"stderr message"}}"#;
let custom_message = r#"{"method":"notifications/custom","params":{"data":"custom"}}"#;
let standard_message =
r#"{"method":"notifications/message","params":{"level":"info","data":"standard"}}"#;
let progress_message = r#"{"method":"notifications/progress","params":{"progressToken":"token","progress":50}}"#;
// Test with valid JSON - all should parse successfully
let result1 =
try_parse_with_compatibility::<serde_json::Value>(stderr_message.as_bytes(), "test");
let result2 =
try_parse_with_compatibility::<serde_json::Value>(custom_message.as_bytes(), "test");
let result3 =
try_parse_with_compatibility::<serde_json::Value>(standard_message.as_bytes(), "test");
let result4 =
try_parse_with_compatibility::<serde_json::Value>(progress_message.as_bytes(), "test");
// All should parse successfully since they're valid JSON
assert!(result1.is_ok());
assert!(result2.is_ok());
assert!(result3.is_ok());
assert!(result4.is_ok());
// Standard notifications should return Some(value)
assert!(result3.unwrap().is_some());
assert!(result4.unwrap().is_some());
println!("Standard notifications are preserved, non-standard are handled gracefully");
}
}
codex-rs/vendor/rmcp/src/transport/auth.rs
Vendored
+1974
View File
File diff suppressed because it is too large Load Diff
codex-rs/vendor/rmcp/src/transport/child_process.rs
Vendored
+309
View File
@@ -0,0 +1,309 @@
use std::process::Stdio;
use futures::future::Future;
use process_wrap::tokio::{ChildWrapper, CommandWrap};
use tokio::{
io::AsyncRead,
process::{ChildStderr, ChildStdin, ChildStdout},
};
use super::{RxJsonRpcMessage, Transport, TxJsonRpcMessage, async_rw::AsyncRwTransport};
use crate::RoleClient;
const MAX_WAIT_ON_DROP_SECS: u64 = 3;
/// The parts of a child process.
type ChildProcessParts = (
Box<dyn ChildWrapper>,
ChildStdout,
ChildStdin,
Option<ChildStderr>,
);
/// Extract the stdio handles from a spawned child.
/// Returns `(child, stdout, stdin, stderr)` where `stderr` is `Some` only
/// if the process was spawned with `Stdio::piped()`.
#[inline]
fn child_process(mut child: Box<dyn ChildWrapper>) -> std::io::Result<ChildProcessParts> {
let child_stdin = match child.inner_mut().stdin().take() {
Some(stdin) => stdin,
None => return Err(std::io::Error::other("stdin was already taken")),
};
let child_stdout = match child.inner_mut().stdout().take() {
Some(stdout) => stdout,
None => return Err(std::io::Error::other("stdout was already taken")),
};
let child_stderr = child.inner_mut().stderr().take();
Ok((child, child_stdout, child_stdin, child_stderr))
}
pub struct TokioChildProcess {
child: ChildWithCleanup,
transport: AsyncRwTransport<RoleClient, ChildStdout, ChildStdin>,
}
pub struct ChildWithCleanup {
inner: Option<Box<dyn ChildWrapper>>,
}
impl Drop for ChildWithCleanup {
fn drop(&mut self) {
// We should not use start_kill(), instead we should use kill() to avoid zombies
if let Some(mut inner) = self.inner.take() {
// We don't care about the result, just try to kill it
tokio::spawn(async move {
if let Err(e) = Box::into_pin(inner.kill()).await {
tracing::warn!("Error killing child process: {}", e);
}
});
}
}
}
// we hold the child process with stdout, for it's easier to implement AsyncRead
pin_project_lite::pin_project! {
pub struct TokioChildProcessOut {
child: ChildWithCleanup,
#[pin]
child_stdout: ChildStdout,
}
}
impl TokioChildProcessOut {
/// Get the process ID of the child process.
pub fn id(&self) -> Option<u32> {
self.child.inner.as_ref()?.id()
}
}
impl AsyncRead for TokioChildProcessOut {
fn poll_read(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
self.project().child_stdout.poll_read(cx, buf)
}
}
impl TokioChildProcess {
/// Convenience: spawn with default `piped` stdio
pub fn new(command: impl Into<CommandWrap>) -> std::io::Result<Self> {
let (proc, _ignored) = TokioChildProcessBuilder::new(command).spawn()?;
Ok(proc)
}
/// Builder entry-point allowing fine-grained stdio control.
pub fn builder(command: impl Into<CommandWrap>) -> TokioChildProcessBuilder {
TokioChildProcessBuilder::new(command)
}
/// Get the process ID of the child process.
pub fn id(&self) -> Option<u32> {
self.child.inner.as_ref()?.id()
}
/// Gracefully shutdown the child process
///
/// This will first close the transport to the child process (the server),
/// and wait for the child process to exit normally with a timeout.
/// If the child process doesn't exit within the timeout, it will be killed.
pub async fn graceful_shutdown(&mut self) -> std::io::Result<()> {
if let Some(mut child) = self.child.inner.take() {
self.transport.close().await?;
let wait_fut = child.wait();
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_secs(MAX_WAIT_ON_DROP_SECS)) => {
if let Err(e) = Box::into_pin(child.kill()).await {
tracing::warn!("Error killing child: {e}");
return Err(e);
}
},
res = wait_fut => {
match res {
Ok(status) => {
tracing::info!("Child exited gracefully {}", status);
}
Err(e) => {
tracing::warn!("Error waiting for child: {e}");
return Err(e);
}
}
}
}
}
Ok(())
}
/// Take ownership of the inner child process
pub fn into_inner(mut self) -> Option<Box<dyn ChildWrapper>> {
self.child.inner.take()
}
/// Split this helper into a reader (stdout) and writer (stdin).
#[deprecated(
since = "0.5.0",
note = "use the Transport trait implementation instead"
)]
pub fn split(self) -> (TokioChildProcessOut, ChildStdin) {
unimplemented!("This method is deprecated, use the Transport trait implementation instead");
}
}
/// Builder for `TokioChildProcess` allowing custom `Stdio` configuration.
pub struct TokioChildProcessBuilder {
cmd: CommandWrap,
stdin: Stdio,
stdout: Stdio,
stderr: Stdio,
}
impl TokioChildProcessBuilder {
fn new(cmd: impl Into<CommandWrap>) -> Self {
Self {
cmd: cmd.into(),
stdin: Stdio::piped(),
stdout: Stdio::piped(),
stderr: Stdio::inherit(),
}
}
/// Override the child stdin configuration.
pub fn stdin(mut self, io: impl Into<Stdio>) -> Self {
self.stdin = io.into();
self
}
/// Override the child stdout configuration.
pub fn stdout(mut self, io: impl Into<Stdio>) -> Self {
self.stdout = io.into();
self
}
/// Override the child stderr configuration.
pub fn stderr(mut self, io: impl Into<Stdio>) -> Self {
self.stderr = io.into();
self
}
/// Spawn the child process. Returns the transport plus an optional captured stderr handle.
pub fn spawn(mut self) -> std::io::Result<(TokioChildProcess, Option<ChildStderr>)> {
self.cmd
.command_mut()
.stdin(self.stdin)
.stdout(self.stdout)
.stderr(self.stderr);
let (child, stdout, stdin, stderr_opt) = child_process(self.cmd.spawn()?)?;
let transport = AsyncRwTransport::new(stdout, stdin);
let proc = TokioChildProcess {
child: ChildWithCleanup { inner: Some(child) },
transport,
};
Ok((proc, stderr_opt))
}
}
impl Transport<RoleClient> for TokioChildProcess {
type Error = std::io::Error;
fn send(
&mut self,
item: TxJsonRpcMessage<RoleClient>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
self.transport.send(item)
}
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<RoleClient>>> + Send {
self.transport.receive()
}
fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send {
self.graceful_shutdown()
}
}
pub trait ConfigureCommandExt {
fn configure(self, f: impl FnOnce(&mut Self)) -> Self;
}
impl ConfigureCommandExt for tokio::process::Command {
fn configure(mut self, f: impl FnOnce(&mut Self)) -> Self {
f(&mut self);
self
}
}
#[cfg(unix)]
#[cfg(test)]
mod tests {
use tokio::process::Command;
use super::*;
#[tokio::test]
async fn test_tokio_child_process_drop() {
let r = TokioChildProcess::new(Command::new("sleep").configure(|cmd| {
cmd.arg("30");
}));
assert!(r.is_ok());
let child_process = r.unwrap();
let id = child_process.id();
assert!(id.is_some());
let id = id.unwrap();
// Drop the child process
drop(child_process);
// Wait a moment to allow the cleanup task to run
tokio::time::sleep(std::time::Duration::from_secs(MAX_WAIT_ON_DROP_SECS + 1)).await;
// Check if the process is still running
let status = Command::new("ps")
.arg("-p")
.arg(id.to_string())
.status()
.await;
match status {
Ok(status) => {
assert!(
!status.success(),
"Process with PID {} is still running",
id
);
}
Err(e) => {
panic!("Failed to check process status: {}", e);
}
}
}
#[tokio::test]
async fn test_tokio_child_process_graceful_shutdown() {
let r = TokioChildProcess::new(Command::new("sleep").configure(|cmd| {
cmd.arg("30");
}));
assert!(r.is_ok());
let mut child_process = r.unwrap();
let id = child_process.id();
assert!(id.is_some());
let id = id.unwrap();
child_process.graceful_shutdown().await.unwrap();
// Wait a moment to allow the cleanup task to run
tokio::time::sleep(std::time::Duration::from_secs(MAX_WAIT_ON_DROP_SECS + 1)).await;
// Check if the process is still running
let status = Command::new("ps")
.arg("-p")
.arg(id.to_string())
.status()
.await;
match status {
Ok(status) => {
assert!(
!status.success(),
"Process with PID {} is still running",
id
);
}
Err(e) => {
panic!("Failed to check process status: {}", e);
}
}
}
}
codex-rs/vendor/rmcp/src/transport/common.rs
Vendored
+19
View File
@@ -0,0 +1,19 @@
#[cfg(feature = "transport-streamable-http-server")]
pub mod server_side_http;
pub mod http_header;
#[cfg(feature = "__reqwest")]
#[cfg_attr(docsrs, doc(cfg(feature = "reqwest")))]
mod reqwest;
// Note: This module provides SSE stream parsing and auto-reconnect utilities.
// It's used by the streamable HTTP client (which receives SSE-formatted responses),
// not the removed SSE transport. The name is historical.
#[cfg(feature = "client-side-sse")]
#[cfg_attr(docsrs, doc(cfg(feature = "client-side-sse")))]
pub mod client_side_sse;
#[cfg(feature = "auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
pub mod auth;
codex-rs/vendor/rmcp/src/transport/common/auth.rs
Vendored
+3
View File
@@ -0,0 +1,3 @@
#[cfg(feature = "transport-streamable-http-client")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client")))]
mod streamable_http_client;
codex-rs/vendor/rmcp/src/transport/common/auth/streamable_http_client.rs
Vendored
+61
View File
@@ -0,0 +1,61 @@
use crate::transport::{
auth::AuthClient,
streamable_http_client::{StreamableHttpClient, StreamableHttpError},
};
impl<C> StreamableHttpClient for AuthClient<C>
where
C: StreamableHttpClient + Send + Sync,
{
type Error = C::Error;
async fn delete_session(
&self,
uri: std::sync::Arc<str>,
session_id: std::sync::Arc<str>,
mut auth_token: Option<String>,
) -> Result<(), crate::transport::streamable_http_client::StreamableHttpError<Self::Error>>
{
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.delete_session(uri, session_id, auth_token)
.await
}
async fn get_stream(
&self,
uri: std::sync::Arc<str>,
session_id: std::sync::Arc<str>,
last_event_id: Option<String>,
mut auth_token: Option<String>,
) -> Result<
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.get_stream(uri, session_id, last_event_id, auth_token)
.await
}
async fn post_message(
&self,
uri: std::sync::Arc<str>,
message: crate::model::ClientJsonRpcMessage,
session_id: Option<std::sync::Arc<str>>,
mut auth_token: Option<String>,
) -> Result<
crate::transport::streamable_http_client::StreamableHttpPostResponse,
StreamableHttpError<Self::Error>,
> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.post_message(uri, message, session_id, auth_token)
.await
}
}
codex-rs/vendor/rmcp/src/transport/common/client_side_sse.rs
Vendored
+310
View File
@@ -0,0 +1,310 @@
use std::{
pin::Pin,
sync::Arc,
task::{Poll, ready},
time::Duration,
};
use futures::{Stream, stream::BoxStream};
use sse_stream::{Error as SseError, Sse};
use crate::model::ServerJsonRpcMessage;
pub type BoxedSseResponse = BoxStream<'static, Result<Sse, SseError>>;
pub trait SseRetryPolicy: std::fmt::Debug + Send + Sync {
fn retry(&self, current_times: usize) -> Option<Duration>;
}
#[derive(Debug, Clone)]
pub struct FixedInterval {
pub max_times: Option<usize>,
pub duration: Duration,
}
impl SseRetryPolicy for FixedInterval {
fn retry(&self, current_times: usize) -> Option<Duration> {
if let Some(max_times) = self.max_times {
if current_times >= max_times {
return None;
}
}
Some(self.duration)
}
}
impl FixedInterval {
pub const DEFAULT_MIN_DURATION: Duration = Duration::from_millis(1000);
}
impl Default for FixedInterval {
fn default() -> Self {
Self {
max_times: None,
duration: Self::DEFAULT_MIN_DURATION,
}
}
}
#[derive(Debug, Clone)]
pub struct ExponentialBackoff {
pub max_times: Option<usize>,
pub base_duration: Duration,
}
impl ExponentialBackoff {
pub const DEFAULT_DURATION: Duration = Duration::from_millis(1000);
}
impl Default for ExponentialBackoff {
fn default() -> Self {
Self {
max_times: None,
base_duration: Self::DEFAULT_DURATION,
}
}
}
impl SseRetryPolicy for ExponentialBackoff {
fn retry(&self, current_times: usize) -> Option<Duration> {
if let Some(max_times) = self.max_times {
if current_times >= max_times {
return None;
}
}
Some(self.base_duration * (2u32.pow(current_times as u32)))
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct NeverRetry;
impl SseRetryPolicy for NeverRetry {
fn retry(&self, _current_times: usize) -> Option<Duration> {
None
}
}
#[derive(Debug, Default)]
pub struct NeverReconnect<E> {
error: Option<E>,
}
impl<E: std::error::Error + Send> SseStreamReconnect for NeverReconnect<E> {
type Error = E;
type Future = futures::future::Ready<Result<BoxedSseResponse, Self::Error>>;
fn retry_connection(&mut self, _last_event_id: Option<&str>) -> Self::Future {
futures::future::ready(Err(self.error.take().expect("should not be called again")))
}
}
/// Abstraction for SSE reconnection logic. Implementors can hook into
/// [`handle_control_event`](Self::handle_control_event) to consume control
/// frames (e.g. `event: endpoint`) that arrive when a server restarts an SSE
/// stream. The default implementation is a no-op, keeping existing behaviour
/// intact.
pub(crate) trait SseStreamReconnect {
type Error: std::error::Error;
type Future: Future<Output = Result<BoxedSseResponse, Self::Error>> + Send;
fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future;
fn handle_control_event(&mut self, _event: &Sse) -> Result<(), Self::Error> {
Ok(())
}
fn handle_stream_error(
&mut self,
error: &(dyn std::error::Error + 'static),
last_event_id: Option<&str>,
) {
if let Some(id) = last_event_id {
tracing::warn!(%id, "sse stream error: {error}");
} else {
tracing::warn!("sse stream error: {error}");
}
}
}
pin_project_lite::pin_project! {
pub(crate) struct SseAutoReconnectStream<R>
where R: SseStreamReconnect
{
retry_policy: Arc<dyn SseRetryPolicy>,
last_event_id: Option<String>,
server_retry_interval: Option<Duration>,
connector: R,
#[pin]
state: SseAutoReconnectStreamState<R::Future>,
}
}
impl<R: SseStreamReconnect> SseAutoReconnectStream<R> {
pub fn new(
stream: BoxedSseResponse,
connector: R,
retry_policy: Arc<dyn SseRetryPolicy>,
) -> Self {
Self {
retry_policy,
last_event_id: None,
server_retry_interval: None,
connector,
state: SseAutoReconnectStreamState::Connected { stream },
}
}
}
impl<E: std::error::Error + Send> SseAutoReconnectStream<NeverReconnect<E>> {
#[allow(dead_code)]
pub(crate) fn never_reconnect(stream: BoxedSseResponse, error_when_reconnect: E) -> Self {
Self {
retry_policy: Arc::new(NeverRetry),
last_event_id: None,
server_retry_interval: None,
connector: NeverReconnect {
error: Some(error_when_reconnect),
},
state: SseAutoReconnectStreamState::Connected { stream },
}
}
}
pin_project_lite::pin_project! {
#[project = SseAutoReconnectStreamStateProj]
pub enum SseAutoReconnectStreamState<F> {
Connected {
#[pin]
stream: BoxedSseResponse,
},
Retrying {
retry_times: usize,
#[pin]
retrying: F,
},
WaitingNextRetry {
#[pin]
sleep: tokio::time::Sleep,
retry_times: usize,
},
Terminated,
}
}
impl<R> Stream for SseAutoReconnectStream<R>
where
R: SseStreamReconnect,
{
type Item = Result<ServerJsonRpcMessage, R::Error>;
fn poll_next(
mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
let mut this = self.as_mut().project();
// let this_state = this.state.as_mut().project()
let state = this.state.as_mut().project();
let next_state = match state {
SseAutoReconnectStreamStateProj::Connected { stream } => {
match ready!(stream.poll_next(cx)) {
Some(Ok(sse)) => {
if let Some(new_server_retry) = sse.retry {
*this.server_retry_interval =
Some(Duration::from_millis(new_server_retry));
}
if let Some(ref event_id) = sse.id {
*this.last_event_id = Some(event_id.clone());
}
// Only treat blank/`message` events as JSON-RPC payloads.
// Other control frames (endpoint, ping, etc.) are passed to
// the reconnection handler.
let is_message_event =
matches!(sse.event.as_deref(), None | Some("") | Some("message"));
if !is_message_event {
match this.connector.handle_control_event(&sse) {
Ok(()) => return self.poll_next(cx),
Err(e) => {
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(Some(Err(e)));
}
}
}
if let Some(data) = sse.data {
match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
Err(e) => {
// Downgrade to debug to avoid noisy logs when servers emit
// non-JSON payloads as message frames. Include last_event_id
// to aid troubleshooting while keeping default behaviour.
let last_id = this.last_event_id.as_deref().unwrap_or("");
tracing::debug!(last_event_id=%last_id, "failed to deserialize server message: {e}");
return self.poll_next(cx);
}
Ok(message) => {
return Poll::Ready(Some(Ok(message)));
}
};
} else {
return self.poll_next(cx);
}
}
Some(Err(e)) => {
this.connector
.handle_stream_error(&e, this.last_event_id.as_deref());
let retrying = this
.connector
.retry_connection(this.last_event_id.as_deref());
SseAutoReconnectStreamState::Retrying {
retry_times: 0,
retrying,
}
}
None => {
tracing::debug!("sse stream terminated");
return Poll::Ready(None);
}
}
}
SseAutoReconnectStreamStateProj::Retrying {
retry_times,
retrying,
} => {
let retry_result = ready!(retrying.poll(cx));
match retry_result {
Ok(new_stream) => SseAutoReconnectStreamState::Connected { stream: new_stream },
Err(e) => {
tracing::debug!("retry sse stream error: {e}");
*retry_times += 1;
if let Some(interval) = this.retry_policy.retry(*retry_times) {
let interval = this
.server_retry_interval
.map(|server_retry_interval| server_retry_interval.max(interval))
.unwrap_or(interval);
let sleep = tokio::time::sleep(interval);
SseAutoReconnectStreamState::WaitingNextRetry {
sleep,
retry_times: *retry_times,
}
} else {
tracing::error!("sse stream error: {e}, max retry times reached");
this.state.set(SseAutoReconnectStreamState::Terminated);
return Poll::Ready(Some(Err(e)));
}
}
}
}
SseAutoReconnectStreamStateProj::WaitingNextRetry { sleep, retry_times } => {
ready!(sleep.poll(cx));
let retrying = this
.connector
.retry_connection(this.last_event_id.as_deref());
let retry_times = *retry_times;
SseAutoReconnectStreamState::Retrying {
retry_times,
retrying,
}
}
SseAutoReconnectStreamStateProj::Terminated => {
return Poll::Ready(None);
}
};
// update the state
this.state.set(next_state);
self.poll_next(cx)
}
}
codex-rs/vendor/rmcp/src/transport/common/http_header.rs
Vendored
+4
View File
@@ -0,0 +1,4 @@
pub const HEADER_SESSION_ID: &str = "Mcp-Session-Id";
pub const HEADER_LAST_EVENT_ID: &str = "Last-Event-Id";
pub const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
pub const JSON_MIME_TYPE: &str = "application/json";
codex-rs/vendor/rmcp/src/transport/common/reqwest.rs
Vendored
+3
View File
@@ -0,0 +1,3 @@
#[cfg(feature = "transport-streamable-http-client-reqwest")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-client-reqwest")))]
mod streamable_http_client;
codex-rs/vendor/rmcp/src/transport/common/reqwest/streamable_http_client.rs
Vendored
+199
View File
@@ -0,0 +1,199 @@
use std::{borrow::Cow, sync::Arc};
use futures::{StreamExt, stream::BoxStream};
use http::header::WWW_AUTHENTICATE;
use reqwest::header::ACCEPT;
use sse_stream::{Sse, SseStream};
use crate::{
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::{
common::http_header::{
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
},
streamable_http_client::*,
},
};
impl From<reqwest::Error> for StreamableHttpError<reqwest::Error> {
fn from(e: reqwest::Error) -> Self {
StreamableHttpError::Client(e)
}
}
impl StreamableHttpClient for reqwest::Client {
type Error = reqwest::Error;
async fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
last_event_id: Option<String>,
auth_token: Option<String>,
) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
let mut request_builder = self
.get(uri.as_ref())
.header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "))
.header(HEADER_SESSION_ID, session_id.as_ref());
if let Some(last_event_id) = last_event_id {
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
}
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
let response = request_builder.send().await?;
if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
return Err(StreamableHttpError::ServerDoesNotSupportSse);
}
let response = response.error_for_status()?;
match response.headers().get(reqwest::header::CONTENT_TYPE) {
Some(ct) => {
if !ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes())
&& !ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes())
{
return Err(StreamableHttpError::UnexpectedContentType(Some(
String::from_utf8_lossy(ct.as_bytes()).to_string(),
)));
}
}
None => {
return Err(StreamableHttpError::UnexpectedContentType(None));
}
}
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(event_stream)
}
async fn delete_session(
&self,
uri: Arc<str>,
session: Arc<str>,
auth_token: Option<String>,
) -> Result<(), StreamableHttpError<Self::Error>> {
let mut request_builder = self.delete(uri.as_ref());
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
let response = request_builder
.header(HEADER_SESSION_ID, session.as_ref())
.send()
.await?;
// if method no allowed
if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
tracing::debug!("this server doesn't support deleting session");
return Ok(());
}
let _response = response.error_for_status()?;
Ok(())
}
async fn post_message(
&self,
uri: Arc<str>,
message: ClientJsonRpcMessage,
session_id: Option<Arc<str>>,
auth_token: Option<String>,
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
let mut request = self
.post(uri.as_ref())
.header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "));
if let Some(auth_header) = auth_token {
request = request.bearer_auth(auth_header);
}
if let Some(session_id) = session_id {
request = request.header(HEADER_SESSION_ID, session_id.as_ref());
}
let response = request.json(&message).send().await?;
if response.status() == reqwest::StatusCode::UNAUTHORIZED {
if let Some(header) = response.headers().get(WWW_AUTHENTICATE) {
let header = header
.to_str()
.map_err(|_| {
StreamableHttpError::UnexpectedServerResponse(Cow::from(
"invalid www-authenticate header value",
))
})?
.to_string();
return Err(StreamableHttpError::AuthRequired(AuthRequiredError {
www_authenticate_header: header,
}));
}
}
let status = response.status();
if matches!(
status,
reqwest::StatusCode::ACCEPTED | reqwest::StatusCode::NO_CONTENT
) {
return Ok(StreamableHttpPostResponse::Accepted);
}
let content_type = response.headers().get(reqwest::header::CONTENT_TYPE);
let session_id = response.headers().get(HEADER_SESSION_ID);
let session_id = session_id
.and_then(|v| v.to_str().ok())
.map(|s| s.to_string());
match content_type {
Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => {
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(StreamableHttpPostResponse::Sse(event_stream, session_id))
}
Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {
let message: ServerJsonRpcMessage = response.json().await?;
Ok(StreamableHttpPostResponse::Json(message, session_id))
}
_ => {
// unexpected content type
tracing::error!("unexpected content type: {:?}", content_type);
Err(StreamableHttpError::UnexpectedContentType(
content_type.map(|ct| String::from_utf8_lossy(ct.as_bytes()).to_string()),
))
}
}
}
}
impl StreamableHttpClientTransport<reqwest::Client> {
/// Creates a new transport using reqwest with the specified URI.
///
/// This is a convenience method that creates a transport using the default
/// reqwest client. This method is only available when the
/// `transport-streamable-http-client-reqwest` feature is enabled.
///
/// # Arguments
///
/// * `uri` - The server URI to connect to
///
/// # Example
///
/// ```rust,no_run
/// use rmcp::transport::StreamableHttpClientTransport;
///
/// // Enable the reqwest feature in Cargo.toml:
/// // rmcp = { version = "0.5", features = ["transport-streamable-http-client-reqwest"] }
///
/// let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp");
/// ```
///
/// # Feature requirement
///
/// This method requires the `transport-streamable-http-client-reqwest` feature.
pub fn from_uri(uri: impl Into<Arc<str>>) -> Self {
StreamableHttpClientTransport::with_client(
reqwest::Client::default(),
StreamableHttpClientTransportConfig {
uri: uri.into(),
auth_header: None,
..Default::default()
},
)
}
/// Build this transport form a config
///
/// # Arguments
///
/// * `config` - The config to use with this transport
pub fn from_config(config: StreamableHttpClientTransportConfig) -> Self {
StreamableHttpClientTransport::with_client(reqwest::Client::default(), config)
}
}
codex-rs/vendor/rmcp/src/transport/common/server_side_http.rs
Vendored
+170
View File
@@ -0,0 +1,170 @@
#![allow(dead_code)]
use std::{convert::Infallible, fmt::Display, sync::Arc, time::Duration};
use bytes::{Buf, Bytes};
use http::Response;
use http_body::Body;
use http_body_util::{BodyExt, Empty, Full, combinators::BoxBody};
use sse_stream::{KeepAlive, Sse, SseBody};
use tokio_util::sync::CancellationToken;
use super::http_header::EVENT_STREAM_MIME_TYPE;
use crate::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
pub type SessionId = Arc<str>;
pub fn session_id() -> SessionId {
uuid::Uuid::new_v4().to_string().into()
}
pub const DEFAULT_AUTO_PING_INTERVAL: Duration = Duration::from_secs(15);
pub(crate) type BoxResponse = Response<BoxBody<Bytes, Infallible>>;
pub(crate) fn accepted_response() -> Response<BoxBody<Bytes, Infallible>> {
Response::builder()
.status(http::StatusCode::ACCEPTED)
.body(Empty::new().boxed())
.expect("valid response")
}
pin_project_lite::pin_project! {
struct TokioTimer {
#[pin]
sleep: tokio::time::Sleep,
}
}
impl Future for TokioTimer {
type Output = ();
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
let this = self.project();
this.sleep.poll(cx)
}
}
impl sse_stream::Timer for TokioTimer {
fn from_duration(duration: Duration) -> Self {
Self {
sleep: tokio::time::sleep(duration),
}
}
fn reset(self: std::pin::Pin<&mut Self>, when: std::time::Instant) {
let this = self.project();
this.sleep.reset(tokio::time::Instant::from_std(when));
}
}
#[derive(Debug, Clone)]
pub struct ServerSseMessage {
/// The event ID for this message. When set, clients can use this ID
/// with the `Last-Event-ID` header to resume the stream from this point.
pub event_id: Option<String>,
/// The JSON-RPC message content. Set to `None` for priming events.
/// See [SEP-1699](https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1699)
pub message: Option<Arc<ServerJsonRpcMessage>>,
/// The retry interval hint for clients. Clients should wait this duration
/// before attempting to reconnect. This maps to the SSE `retry:` field.
pub retry: Option<Duration>,
}
pub(crate) fn sse_stream_response(
stream: impl futures::Stream<Item = ServerSseMessage> + Send + Sync + 'static,
keep_alive: Option<Duration>,
ct: CancellationToken,
) -> Response<BoxBody<Bytes, Infallible>> {
use futures::StreamExt;
let stream = stream
.map(|message| {
let mut sse = if let Some(ref msg) = message.message {
let data = serde_json::to_string(msg.as_ref()).expect("valid message");
Sse::default().data(data)
} else {
// Priming event: empty data per SEP-1699 (just "data:\n")
Sse::default().data("")
};
sse.id = message.event_id;
if let Some(retry) = message.retry {
sse.retry = Some(retry.as_millis() as u64);
}
Result::<Sse, Infallible>::Ok(sse)
})
.take_until(async move { ct.cancelled().await });
let stream = SseBody::new(stream);
let stream = match keep_alive {
Some(duration) => stream
.with_keep_alive::<TokioTimer>(KeepAlive::new().interval(duration))
.boxed(),
None => stream.boxed(),
};
Response::builder()
.status(http::StatusCode::OK)
.header(http::header::CONTENT_TYPE, EVENT_STREAM_MIME_TYPE)
.header(http::header::CACHE_CONTROL, "no-cache")
.body(stream)
.expect("valid response")
}
pub(crate) const fn internal_error_response<E: Display>(
context: &str,
) -> impl FnOnce(E) -> Response<BoxBody<Bytes, Infallible>> {
move |error| {
tracing::error!("Internal server error when {context}: {error}");
Response::builder()
.status(http::StatusCode::INTERNAL_SERVER_ERROR)
.body(
Full::new(Bytes::from(format!(
"Encounter an error when {context}: {error}"
)))
.boxed(),
)
.expect("valid response")
}
}
pub(crate) fn unexpected_message_response(expect: &str) -> Response<BoxBody<Bytes, Infallible>> {
Response::builder()
.status(http::StatusCode::UNPROCESSABLE_ENTITY)
.body(Full::new(Bytes::from(format!("Unexpected message, expect {expect}"))).boxed())
.expect("valid response")
}
pub(crate) async fn expect_json<B>(
body: B,
) -> Result<ClientJsonRpcMessage, Response<BoxBody<Bytes, Infallible>>>
where
B: Body + Send + 'static,
B::Error: Display,
{
match body.collect().await {
Ok(bytes) => {
match serde_json::from_reader::<_, ClientJsonRpcMessage>(bytes.aggregate().reader()) {
Ok(message) => Ok(message),
Err(e) => {
let response = Response::builder()
.status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
.body(
Full::new(Bytes::from(format!("fail to deserialize request body {e}")))
.boxed(),
)
.expect("valid response");
Err(response)
}
}
}
Err(e) => {
let response = Response::builder()
.status(http::StatusCode::INTERNAL_SERVER_ERROR)
.body(Full::new(Bytes::from(format!("Failed to read request body: {e}"))).boxed())
.expect("valid response");
Err(response)
}
}
}
codex-rs/vendor/rmcp/src/transport/io.rs
Vendored
+6
View File
@@ -0,0 +1,6 @@
/// # StdIO Transport
///
/// Create a pair of [`tokio::io::Stdin`] and [`tokio::io::Stdout`].
pub fn stdio() -> (tokio::io::Stdin, tokio::io::Stdout) {
(tokio::io::stdin(), tokio::io::stdout())
}
codex-rs/vendor/rmcp/src/transport/sink_stream.rs
Vendored
+78
View File
@@ -0,0 +1,78 @@
use std::sync::Arc;
use futures::{Sink, Stream};
use tokio::sync::Mutex;
use super::{IntoTransport, Transport};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
pub struct SinkStreamTransport<Si, St> {
stream: St,
sink: Arc<Mutex<Si>>,
}
impl<Si, St> SinkStreamTransport<Si, St> {
pub fn new(sink: Si, stream: St) -> Self {
Self {
stream,
sink: Arc::new(Mutex::new(sink)),
}
}
}
impl<Role: ServiceRole, Si, St> Transport<Role> for SinkStreamTransport<Si, St>
where
St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin,
Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
Si::Error: std::error::Error + Send + Sync + 'static,
{
type Error = Si::Error;
fn send(
&mut self,
item: TxJsonRpcMessage<Role>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
use futures::SinkExt;
let lock = self.sink.clone();
async move {
let mut write = lock.lock().await;
write.send(item).await
}
}
fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<Role>>> {
use futures::StreamExt;
self.stream.next()
}
async fn close(&mut self) -> Result<(), Self::Error> {
Ok(())
}
}
pub enum TransportAdapterSinkStream {}
impl<Role, Si, St> IntoTransport<Role, Si::Error, TransportAdapterSinkStream> for (Si, St)
where
Role: ServiceRole,
Si: Send + Sink<TxJsonRpcMessage<Role>> + Unpin + 'static,
St: Send + Stream<Item = RxJsonRpcMessage<Role>> + Unpin + 'static,
Si::Error: std::error::Error + Send + Sync + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = Si::Error> + 'static {
SinkStreamTransport::new(self.0, self.1)
}
}
pub enum TransportAdapterAsyncCombinedRW {}
impl<Role, S> IntoTransport<Role, S::Error, TransportAdapterAsyncCombinedRW> for S
where
Role: ServiceRole,
S: Sink<TxJsonRpcMessage<Role>> + Stream<Item = RxJsonRpcMessage<Role>> + Send + 'static,
S::Error: std::error::Error + Send + Sync + 'static,
{
fn into_transport(self) -> impl Transport<Role, Error = S::Error> + 'static {
use futures::StreamExt;
IntoTransport::<Role, S::Error, TransportAdapterSinkStream>::into_transport(self.split())
}
}
codex-rs/vendor/rmcp/src/transport/streamable_http_client.rs
Vendored
+774
View File
@@ -0,0 +1,774 @@
use std::{borrow::Cow, sync::Arc, time::Duration};
use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream};
pub use sse_stream::Error as SseError;
use sse_stream::Sse;
use thiserror::Error;
use tokio_util::sync::CancellationToken;
use tracing::debug;
use super::common::client_side_sse::{ExponentialBackoff, SseRetryPolicy, SseStreamReconnect};
use crate::{
RoleClient,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::{
common::client_side_sse::SseAutoReconnectStream,
worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport},
},
};
type BoxedSseStream = BoxStream<'static, Result<Sse, SseError>>;
#[derive(Debug)]
pub struct AuthRequiredError {
pub www_authenticate_header: String,
}
#[derive(Error, Debug)]
pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
#[error("SSE error: {0}")]
Sse(#[from] SseError),
#[error("Io error: {0}")]
Io(#[from] std::io::Error),
#[error("Client error: {0}")]
Client(E),
#[error("unexpected end of stream")]
UnexpectedEndOfStream,
#[error("unexpected server response: {0}")]
UnexpectedServerResponse(Cow<'static, str>),
#[error("Unexpected content type: {0:?}")]
UnexpectedContentType(Option<String>),
#[error("Server does not support SSE")]
ServerDoesNotSupportSse,
#[error("Server does not support delete session")]
ServerDoesNotSupportDeleteSession,
#[error("Tokio join error: {0}")]
TokioJoinError(#[from] tokio::task::JoinError),
#[error("Deserialize error: {0}")]
Deserialize(#[from] serde_json::Error),
#[error("Transport channel closed")]
TransportChannelClosed,
#[error("Missing session id in HTTP response")]
MissingSessionIdInResponse,
#[cfg(feature = "auth")]
#[cfg_attr(docsrs, doc(cfg(feature = "auth")))]
#[error("Auth error: {0}")]
Auth(#[from] crate::transport::auth::AuthError),
#[error("Auth required")]
AuthRequired(AuthRequiredError),
}
#[derive(Debug, Clone, Error)]
pub enum StreamableHttpProtocolError {
#[error("Missing session id in response")]
MissingSessionIdInResponse,
}
#[allow(clippy::large_enum_variant)]
pub enum StreamableHttpPostResponse {
Accepted,
Json(ServerJsonRpcMessage, Option<String>),
Sse(BoxedSseStream, Option<String>),
}
impl std::fmt::Debug for StreamableHttpPostResponse {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Accepted => write!(f, "Accepted"),
Self::Json(arg0, arg1) => f.debug_tuple("Json").field(arg0).field(arg1).finish(),
Self::Sse(_, arg1) => f.debug_tuple("Sse").field(arg1).finish(),
}
}
}
impl StreamableHttpPostResponse {
pub async fn expect_initialized<E>(
self,
) -> Result<(ServerJsonRpcMessage, Option<String>), StreamableHttpError<E>>
where
E: std::error::Error + Send + Sync + 'static,
{
match self {
Self::Json(message, session_id) => Ok((message, session_id)),
Self::Sse(mut stream, session_id) => {
while let Some(event) = stream.next().await {
let event = event?;
let payload = event.data.unwrap_or_default();
if payload.trim().is_empty() {
continue;
}
let message: ServerJsonRpcMessage = serde_json::from_str(&payload)?;
if matches!(message, ServerJsonRpcMessage::Response(_)) {
return Ok((message, session_id));
}
debug!(
?message,
"received message before initialize response; continuing to drain stream"
);
}
Err(StreamableHttpError::UnexpectedServerResponse(
"empty sse stream".into(),
))
}
_ => Err(StreamableHttpError::UnexpectedServerResponse(
"expect initialized, accepted".into(),
)),
}
}
pub fn expect_json<E>(self) -> Result<ServerJsonRpcMessage, StreamableHttpError<E>>
where
E: std::error::Error + Send + Sync + 'static,
{
match self {
Self::Json(message, ..) => Ok(message),
got => Err(StreamableHttpError::UnexpectedServerResponse(
format!("expect json, got {got:?}").into(),
)),
}
}
pub fn expect_accepted<E>(self) -> Result<(), StreamableHttpError<E>>
where
E: std::error::Error + Send + Sync + 'static,
{
match self {
Self::Accepted => Ok(()),
got => Err(StreamableHttpError::UnexpectedServerResponse(
format!("expect accepted, got {got:?}").into(),
)),
}
}
}
pub trait StreamableHttpClient: Clone + Send + 'static {
type Error: std::error::Error + Send + Sync + 'static;
fn post_message(
&self,
uri: Arc<str>,
message: ClientJsonRpcMessage,
session_id: Option<Arc<str>>,
auth_header: Option<String>,
) -> impl Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>>>
+ Send
+ '_;
fn delete_session(
&self,
uri: Arc<str>,
session_id: Arc<str>,
auth_header: Option<String>,
) -> impl Future<Output = Result<(), StreamableHttpError<Self::Error>>> + Send + '_;
fn get_stream(
&self,
uri: Arc<str>,
session_id: Arc<str>,
last_event_id: Option<String>,
auth_header: Option<String>,
) -> impl Future<
Output = Result<
BoxStream<'static, Result<Sse, SseError>>,
StreamableHttpError<Self::Error>,
>,
> + Send
+ '_;
}
pub struct RetryConfig {
pub max_times: Option<usize>,
pub min_duration: Duration,
}
struct StreamableHttpClientReconnect<C> {
pub client: C,
pub session_id: Arc<str>,
pub uri: Arc<str>,
pub auth_header: Option<String>,
}
impl<C: StreamableHttpClient> SseStreamReconnect for StreamableHttpClientReconnect<C> {
type Error = StreamableHttpError<C::Error>;
type Future = BoxFuture<'static, Result<BoxedSseStream, Self::Error>>;
fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future {
let client = self.client.clone();
let uri = self.uri.clone();
let session_id = self.session_id.clone();
let auth_header = self.auth_header.clone();
let last_event_id = last_event_id.map(|s| s.to_owned());
Box::pin(async move {
client
.get_stream(uri, session_id, last_event_id, auth_header)
.await
})
}
}
#[derive(Debug, Clone, Default)]
pub struct StreamableHttpClientWorker<C: StreamableHttpClient> {
pub client: C,
pub config: StreamableHttpClientTransportConfig,
}
impl<C: StreamableHttpClient + Default> StreamableHttpClientWorker<C> {
pub fn new_simple(url: impl Into<Arc<str>>) -> Self {
Self {
client: C::default(),
config: StreamableHttpClientTransportConfig {
uri: url.into(),
..Default::default()
},
}
}
}
impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
pub fn new(client: C, config: StreamableHttpClientTransportConfig) -> Self {
Self { client, config }
}
}
impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
async fn execute_sse_stream(
sse_stream: impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>>
+ Send
+ 'static,
sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>,
close_on_response: bool,
ct: CancellationToken,
) -> Result<(), StreamableHttpError<C::Error>> {
let mut sse_stream = std::pin::pin!(sse_stream);
loop {
let message = tokio::select! {
event = sse_stream.next() => {
event
}
_ = ct.cancelled() => {
tracing::debug!("cancelled");
break;
}
};
let Some(message) = message.transpose()? else {
break;
};
let is_response = matches!(message, ServerJsonRpcMessage::Response(_));
let yield_result = sse_worker_tx.send(message).await;
if yield_result.is_err() {
tracing::trace!("streamable http transport worker dropped, exiting");
break;
}
if close_on_response && is_response {
tracing::debug!("got response, closing sse stream");
break;
}
}
Ok(())
}
}
impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
type Role = RoleClient;
type Error = StreamableHttpError<C::Error>;
fn err_closed() -> Self::Error {
StreamableHttpError::TransportChannelClosed
}
fn err_join(e: tokio::task::JoinError) -> Self::Error {
StreamableHttpError::TokioJoinError(e)
}
fn config(&self) -> super::worker::WorkerConfig {
super::worker::WorkerConfig {
name: Some("StreamableHttpClientWorker".into()),
channel_buffer_capacity: self.config.channel_buffer_capacity,
}
}
async fn run(
self,
mut context: super::worker::WorkerContext<Self>,
) -> Result<(), WorkerQuitReason<Self::Error>> {
let channel_buffer_capacity = self.config.channel_buffer_capacity;
let (sse_worker_tx, mut sse_worker_rx) =
tokio::sync::mpsc::channel::<ServerJsonRpcMessage>(channel_buffer_capacity);
let config = self.config.clone();
let transport_task_ct = context.cancellation_token.clone();
let _drop_guard = transport_task_ct.clone().drop_guard();
let WorkerSendRequest {
responder,
message: initialize_request,
} = context.recv_from_handler().await?;
let (message, session_id) = match self
.client
.post_message(
config.uri.clone(),
initialize_request,
None,
self.config.auth_header,
)
.await
{
Ok(res) => {
let _ = responder.send(Ok(()));
res.expect_initialized::<C::Error>().await.map_err(
WorkerQuitReason::fatal_context("process initialize response"),
)?
}
Err(err) => {
let msg = format!("{:?}", err);
let _ = responder.send(Err(err));
return Err(WorkerQuitReason::fatal(
StreamableHttpError::TransportChannelClosed,
msg,
));
}
};
let session_id: Option<Arc<str>> = if let Some(session_id) = session_id {
Some(session_id.into())
} else {
if !self.config.allow_stateless {
return Err(WorkerQuitReason::fatal(
StreamableHttpError::<C::Error>::MissingSessionIdInResponse,
"process initialize response",
));
}
None
};
// Store session info for cleanup when run() exits (not spawned, so cleanup completes before close() returns)
let session_cleanup_info = session_id.as_ref().map(|sid| {
(
self.client.clone(),
config.uri.clone(),
sid.clone(),
config.auth_header.clone(),
)
});
context.send_to_handler(message).await?;
let initialized_notification = context.recv_from_handler().await?;
// expect a initialized response
self.client
.post_message(
config.uri.clone(),
initialized_notification.message,
session_id.clone(),
config.auth_header.clone(),
)
.await
.map_err(WorkerQuitReason::fatal_context(
"send initialized notification",
))?
.expect_accepted::<C::Error>()
.map_err(WorkerQuitReason::fatal_context(
"process initialized notification response",
))?;
let _ = initialized_notification.responder.send(Ok(()));
#[allow(clippy::large_enum_variant)]
enum Event<W: Worker, E: std::error::Error + Send + Sync + 'static> {
ClientMessage(WorkerSendRequest<W>),
ServerMessage(ServerJsonRpcMessage),
StreamResult(Result<(), StreamableHttpError<E>>),
}
let mut streams = tokio::task::JoinSet::new();
if let Some(session_id) = &session_id {
let client = self.client.clone();
let uri = config.uri.clone();
let session_id = session_id.clone();
let auth_header = config.auth_header.clone();
let retry_config = self.config.retry_config.clone();
let sse_worker_tx = sse_worker_tx.clone();
let transport_task_ct = transport_task_ct.clone();
let config_uri = config.uri.clone();
let config_auth_header = config.auth_header.clone();
streams.spawn(async move {
match client
.get_stream(uri.clone(), session_id.clone(), None, auth_header.clone())
.await
{
Ok(stream) => {
let sse_stream = SseAutoReconnectStream::new(
stream,
StreamableHttpClientReconnect {
client: client.clone(),
session_id: session_id.clone(),
uri: config_uri,
auth_header: config_auth_header,
},
retry_config,
);
Self::execute_sse_stream(
sse_stream,
sse_worker_tx,
false,
transport_task_ct.child_token(),
)
.await
}
Err(StreamableHttpError::ServerDoesNotSupportSse) => {
tracing::debug!("server doesn't support sse, skip common stream");
Ok(())
}
Err(e) => {
// fail to get common stream
tracing::error!("fail to get common stream: {e}");
Err(e)
}
}
});
}
// Main event loop - capture exit reason so we can do cleanup before returning
let loop_result: Result<(), WorkerQuitReason<Self::Error>> = 'main_loop: loop {
let event = tokio::select! {
_ = transport_task_ct.cancelled() => {
tracing::debug!("cancelled");
break 'main_loop Err(WorkerQuitReason::Cancelled);
}
message = context.recv_from_handler() => {
match message {
Ok(msg) => Event::ClientMessage(msg),
Err(e) => break 'main_loop Err(e),
}
},
message = sse_worker_rx.recv() => {
let Some(message) = message else {
tracing::trace!("transport dropped, exiting");
break 'main_loop Err(WorkerQuitReason::HandlerTerminated);
};
Event::ServerMessage(message)
},
terminated_stream = streams.join_next(), if !streams.is_empty() => {
match terminated_stream {
Some(result) => {
Event::StreamResult(result.map_err(StreamableHttpError::TokioJoinError).and_then(std::convert::identity))
}
None => {
continue
}
}
}
};
match event {
Event::ClientMessage(send_request) => {
let WorkerSendRequest { message, responder } = send_request;
let response = self
.client
.post_message(
config.uri.clone(),
message,
session_id.clone(),
config.auth_header.clone(),
)
.await;
let send_result = match response {
Err(e) => Err(e),
Ok(StreamableHttpPostResponse::Accepted) => {
tracing::trace!("client message accepted");
Ok(())
}
Ok(StreamableHttpPostResponse::Json(message, ..)) => {
context.send_to_handler(message).await?;
Ok(())
}
Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
if let Some(session_id) = &session_id {
let sse_stream = SseAutoReconnectStream::new(
stream,
StreamableHttpClientReconnect {
client: self.client.clone(),
session_id: session_id.clone(),
uri: config.uri.clone(),
auth_header: config.auth_header.clone(),
},
self.config.retry_config.clone(),
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
} else {
let sse_stream = SseAutoReconnectStream::never_reconnect(
stream,
StreamableHttpError::<C::Error>::UnexpectedEndOfStream,
);
streams.spawn(Self::execute_sse_stream(
sse_stream,
sse_worker_tx.clone(),
true,
transport_task_ct.child_token(),
));
}
tracing::trace!("got new sse stream");
Ok(())
}
};
let _ = responder.send(send_result);
}
Event::ServerMessage(json_rpc_message) => {
// send the message to the handler
if let Err(e) = context.send_to_handler(json_rpc_message).await {
break 'main_loop Err(e);
}
}
Event::StreamResult(result) => {
if result.is_err() {
tracing::warn!(
"sse client event stream terminated with error: {:?}",
result
);
}
}
}
};
// Cleanup session before returning (ensures close() waits for session deletion)
// Use a timeout to prevent indefinite hangs if the server is unresponsive
if let Some((client, url, session_id, auth_header)) = session_cleanup_info {
const SESSION_CLEANUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
match tokio::time::timeout(
SESSION_CLEANUP_TIMEOUT,
client.delete_session(url, session_id.clone(), auth_header),
)
.await
{
Ok(Ok(_)) => {
tracing::info!(session_id = session_id.as_ref(), "delete session success")
}
Ok(Err(StreamableHttpError::ServerDoesNotSupportDeleteSession)) => {
tracing::info!(
session_id = session_id.as_ref(),
"server doesn't support delete session"
)
}
Ok(Err(e)) => {
tracing::error!(
session_id = session_id.as_ref(),
"fail to delete session: {e}"
);
}
Err(_elapsed) => {
tracing::warn!(
session_id = session_id.as_ref(),
"session cleanup timed out after {:?}",
SESSION_CLEANUP_TIMEOUT
);
}
}
}
loop_result
}
}
/// A client-agnostic HTTP transport for RMCP that supports streaming responses.
///
/// This transport allows you to choose your preferred HTTP client implementation
/// by implementing the [`StreamableHttpClient`] trait. The transport handles
/// session management, SSE streaming, and automatic reconnection.
///
/// # Usage
///
/// ## Using reqwest
///
/// ```rust,no_run
/// use rmcp::transport::StreamableHttpClientTransport;
///
/// // Enable the reqwest feature in Cargo.toml:
/// // rmcp = { version = "0.5", features = ["transport-streamable-http-client-reqwest"] }
///
/// let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp");
/// ```
///
/// ## Using a custom HTTP client
///
/// ```rust,no_run
/// use rmcp::transport::streamable_http_client::{
/// StreamableHttpClient,
/// StreamableHttpClientTransport,
/// StreamableHttpClientTransportConfig
/// };
/// use std::sync::Arc;
/// use futures::stream::BoxStream;
/// use rmcp::model::ClientJsonRpcMessage;
/// use sse_stream::{Sse, Error as SseError};
///
/// #[derive(Clone)]
/// struct MyHttpClient;
///
/// #[derive(Debug, thiserror::Error)]
/// struct MyError;
///
/// impl std::fmt::Display for MyError {
/// fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
/// write!(f, "MyError")
/// }
/// }
///
/// impl StreamableHttpClient for MyHttpClient {
/// type Error = MyError;
///
/// async fn post_message(
/// &self,
/// _uri: Arc<str>,
/// _message: ClientJsonRpcMessage,
/// _session_id: Option<Arc<str>>,
/// _auth_header: Option<String>,
/// ) -> Result<rmcp::transport::streamable_http_client::StreamableHttpPostResponse, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
/// todo!()
/// }
///
/// async fn delete_session(
/// &self,
/// _uri: Arc<str>,
/// _session_id: Arc<str>,
/// _auth_header: Option<String>,
/// ) -> Result<(), rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
/// todo!()
/// }
///
/// async fn get_stream(
/// &self,
/// _uri: Arc<str>,
/// _session_id: Arc<str>,
/// _last_event_id: Option<String>,
/// _auth_header: Option<String>,
/// ) -> Result<BoxStream<'static, Result<Sse, SseError>>, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
/// todo!()
/// }
/// }
///
/// let transport = StreamableHttpClientTransport::with_client(
/// MyHttpClient,
/// StreamableHttpClientTransportConfig::with_uri("http://localhost:8000/mcp")
/// );
/// ```
///
/// # Feature Flags
///
/// - `transport-streamable-http-client`: Base feature providing the generic transport infrastructure
/// - `transport-streamable-http-client-reqwest`: Includes reqwest HTTP client support with convenience methods
pub type StreamableHttpClientTransport<C> = WorkerTransport<StreamableHttpClientWorker<C>>;
impl<C: StreamableHttpClient> StreamableHttpClientTransport<C> {
/// Creates a new transport with a custom HTTP client implementation.
///
/// This method allows you to use any HTTP client that implements the [`StreamableHttpClient`] trait.
/// Use this when you want to use a custom HTTP client or when the reqwest feature is not enabled.
///
/// # Arguments
///
/// * `client` - Your HTTP client implementation
/// * `config` - Transport configuration including the server URI
///
/// # Example
///
/// ```rust,no_run
/// use rmcp::transport::streamable_http_client::{
/// StreamableHttpClient,
/// StreamableHttpClientTransport,
/// StreamableHttpClientTransportConfig
/// };
/// use std::sync::Arc;
/// use futures::stream::BoxStream;
/// use rmcp::model::ClientJsonRpcMessage;
/// use sse_stream::{Sse, Error as SseError};
///
/// // Define your custom client
/// #[derive(Clone)]
/// struct MyHttpClient;
///
/// #[derive(Debug, thiserror::Error)]
/// struct MyError;
///
/// impl std::fmt::Display for MyError {
/// fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
/// write!(f, "MyError")
/// }
/// }
///
/// impl StreamableHttpClient for MyHttpClient {
/// type Error = MyError;
///
/// async fn post_message(
/// &self,
/// _uri: Arc<str>,
/// _message: ClientJsonRpcMessage,
/// _session_id: Option<Arc<str>>,
/// _auth_header: Option<String>,
/// ) -> Result<rmcp::transport::streamable_http_client::StreamableHttpPostResponse, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
/// todo!()
/// }
///
/// async fn delete_session(
/// &self,
/// _uri: Arc<str>,
/// _session_id: Arc<str>,
/// _auth_header: Option<String>,
/// ) -> Result<(), rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
/// todo!()
/// }
///
/// async fn get_stream(
/// &self,
/// _uri: Arc<str>,
/// _session_id: Arc<str>,
/// _last_event_id: Option<String>,
/// _auth_header: Option<String>,
/// ) -> Result<BoxStream<'static, Result<Sse, SseError>>, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
/// todo!()
/// }
/// }
///
/// let transport = StreamableHttpClientTransport::with_client(
/// MyHttpClient,
/// StreamableHttpClientTransportConfig::with_uri("http://localhost:8000/mcp")
/// );
/// ```
pub fn with_client(client: C, config: StreamableHttpClientTransportConfig) -> Self {
let worker = StreamableHttpClientWorker::new(client, config);
WorkerTransport::spawn(worker)
}
}
#[derive(Debug, Clone)]
pub struct StreamableHttpClientTransportConfig {
pub uri: Arc<str>,
pub retry_config: Arc<dyn SseRetryPolicy>,
pub channel_buffer_capacity: usize,
/// if true, the transport will not require a session to be established
pub allow_stateless: bool,
/// The value to send in the authorization header
pub auth_header: Option<String>,
}
impl StreamableHttpClientTransportConfig {
pub fn with_uri(uri: impl Into<Arc<str>>) -> Self {
Self {
uri: uri.into(),
..Default::default()
}
}
/// Set the authorization header to send with requests
///
/// # Arguments
///
/// * `value` - A bearer token without the `Bearer ` prefix
pub fn auth_header<T: Into<String>>(mut self, value: T) -> Self {
// set our authorization header
self.auth_header = Some(value.into());
self
}
}
impl Default for StreamableHttpClientTransportConfig {
fn default() -> Self {
Self {
uri: "localhost".into(),
retry_config: Arc::new(ExponentialBackoff::default()),
channel_buffer_capacity: 16,
allow_stateless: true,
auth_header: None,
}
}
}
codex-rs/vendor/rmcp/src/transport/streamable_http_server.rs
Vendored
+8
View File
@@ -0,0 +1,8 @@
pub mod session;
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub mod tower;
pub use session::{SessionId, SessionManager};
#[cfg(feature = "transport-streamable-http-server")]
#[cfg_attr(docsrs, doc(cfg(feature = "transport-streamable-http-server")))]
pub use tower::{StreamableHttpServerConfig, StreamableHttpService};
codex-rs/vendor/rmcp/src/transport/streamable_http_server/session.rs
Vendored
+53
View File
@@ -0,0 +1,53 @@
use futures::Stream;
pub use crate::transport::common::server_side_http::{ServerSseMessage, SessionId};
use crate::{
RoleServer,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
};
pub mod local;
pub mod never;
pub trait SessionManager: Send + Sync + 'static {
type Error: std::error::Error + Send + 'static;
type Transport: crate::transport::Transport<RoleServer>;
/// Create a new session with the given id and configuration.
fn create_session(
&self,
) -> impl Future<Output = Result<(SessionId, Self::Transport), Self::Error>> + Send;
fn initialize_session(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
) -> impl Future<Output = Result<ServerJsonRpcMessage, Self::Error>> + Send;
fn has_session(&self, id: &SessionId)
-> impl Future<Output = Result<bool, Self::Error>> + Send;
fn close_session(&self, id: &SessionId)
-> impl Future<Output = Result<(), Self::Error>> + Send;
fn create_stream(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + Sync + 'static, Self::Error>,
> + Send;
fn accept_message(
&self,
id: &SessionId,
message: ClientJsonRpcMessage,
) -> impl Future<Output = Result<(), Self::Error>> + Send;
fn create_standalone_stream(
&self,
id: &SessionId,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + Sync + 'static, Self::Error>,
> + Send;
fn resume(
&self,
id: &SessionId,
last_event_id: String,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + Sync + 'static, Self::Error>,
> + Send;
}
codex-rs/vendor/rmcp/src/transport/streamable_http_server/session/local.rs
Vendored
+1043
View File
File diff suppressed because it is too large Load Diff
codex-rs/vendor/rmcp/src/transport/streamable_http_server/session/never.rs
Vendored
+107
View File
@@ -0,0 +1,107 @@
use futures::Stream;
use thiserror::Error;
use super::{ServerSseMessage, SessionId, SessionManager};
use crate::{
RoleServer,
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
transport::Transport,
};
#[derive(Debug, Clone, Error)]
#[error("Session management is not supported")]
pub struct ErrorSessionManagementNotSupported;
#[derive(Debug, Clone, Default)]
pub struct NeverSessionManager {}
pub enum NeverTransport {}
impl Transport<RoleServer> for NeverTransport {
type Error = ErrorSessionManagementNotSupported;
fn send(
&mut self,
_item: ServerJsonRpcMessage,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn receive(&mut self) -> impl Future<Output = Option<ClientJsonRpcMessage>> {
futures::future::ready(None)
}
async fn close(&mut self) -> Result<(), Self::Error> {
Err(ErrorSessionManagementNotSupported)
}
}
impl SessionManager for NeverSessionManager {
type Error = ErrorSessionManagementNotSupported;
type Transport = NeverTransport;
fn create_session(
&self,
) -> impl Future<Output = Result<(SessionId, Self::Transport), Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn initialize_session(
&self,
_id: &SessionId,
_message: ClientJsonRpcMessage,
) -> impl Future<Output = Result<ServerJsonRpcMessage, Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn has_session(
&self,
_id: &SessionId,
) -> impl Future<Output = Result<bool, Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn close_session(
&self,
_id: &SessionId,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
fn create_stream(
&self,
_id: &SessionId,
_message: ClientJsonRpcMessage,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send {
futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
ErrorSessionManagementNotSupported,
))
}
fn create_standalone_stream(
&self,
_id: &SessionId,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send {
futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
ErrorSessionManagementNotSupported,
))
}
fn resume(
&self,
_id: &SessionId,
_last_event_id: String,
) -> impl Future<
Output = Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error>,
> + Send {
futures::future::ready(Result::<futures::stream::Pending<_>, _>::Err(
ErrorSessionManagementNotSupported,
))
}
fn accept_message(
&self,
_id: &SessionId,
_message: ClientJsonRpcMessage,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
futures::future::ready(Err(ErrorSessionManagementNotSupported))
}
}
codex-rs/vendor/rmcp/src/transport/streamable_http_server/tower.rs
Vendored
+521
View File
@@ -0,0 +1,521 @@
use std::{convert::Infallible, fmt::Display, sync::Arc, time::Duration};
use bytes::Bytes;
use futures::{StreamExt, future::BoxFuture};
use http::{Method, Request, Response, header::ALLOW};
use http_body::Body;
use http_body_util::{BodyExt, Full, combinators::BoxBody};
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use super::session::SessionManager;
use crate::{
RoleServer,
model::{ClientJsonRpcMessage, ClientRequest, GetExtensions},
serve_server,
service::serve_directly,
transport::{
OneshotTransport, TransportAdapterIdentity,
common::{
http_header::{
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
},
server_side_http::{
BoxResponse, ServerSseMessage, accepted_response, expect_json,
internal_error_response, sse_stream_response, unexpected_message_response,
},
},
},
};
#[derive(Debug, Clone)]
pub struct StreamableHttpServerConfig {
/// The ping message duration for SSE connections.
pub sse_keep_alive: Option<Duration>,
/// The retry interval for SSE priming events.
pub sse_retry: Option<Duration>,
/// If true, the server will create a session for each request and keep it alive.
/// When enabled, SSE priming events are sent to enable client reconnection.
pub stateful_mode: bool,
/// Cancellation token for the Streamable HTTP server.
///
/// When this token is cancelled, all active sessions are terminated and
/// the server stops accepting new requests.
pub cancellation_token: CancellationToken,
}
impl Default for StreamableHttpServerConfig {
fn default() -> Self {
Self {
sse_keep_alive: Some(Duration::from_secs(15)),
sse_retry: Some(Duration::from_secs(3)),
stateful_mode: true,
cancellation_token: CancellationToken::new(),
}
}
}
/// # Streamable Http Server
///
/// ## Extract information from raw http request
///
/// The http service will consume the request body, however the rest part will be remain and injected into [`crate::model::Extensions`],
/// which you can get from [`crate::service::RequestContext`].
/// ```rust
/// use rmcp::handler::server::tool::Extension;
/// use http::request::Parts;
/// async fn my_tool(Extension(parts): Extension<Parts>) {
/// tracing::info!("http parts:{parts:?}")
/// }
/// ```
pub struct StreamableHttpService<S, M = super::session::local::LocalSessionManager> {
pub config: StreamableHttpServerConfig,
session_manager: Arc<M>,
service_factory: Arc<dyn Fn() -> Result<S, std::io::Error> + Send + Sync>,
}
impl<S, M> Clone for StreamableHttpService<S, M> {
fn clone(&self) -> Self {
Self {
config: self.config.clone(),
session_manager: self.session_manager.clone(),
service_factory: self.service_factory.clone(),
}
}
}
impl<RequestBody, S, M> tower_service::Service<Request<RequestBody>> for StreamableHttpService<S, M>
where
RequestBody: Body + Send + 'static,
S: crate::Service<RoleServer>,
M: SessionManager,
RequestBody::Error: Display,
RequestBody::Data: Send + 'static,
{
type Response = BoxResponse;
type Error = Infallible;
type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
fn call(&mut self, req: http::Request<RequestBody>) -> Self::Future {
let service = self.clone();
Box::pin(async move {
let response = service.handle(req).await;
Ok(response)
})
}
fn poll_ready(
&mut self,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
}
impl<S, M> StreamableHttpService<S, M>
where
S: crate::Service<RoleServer> + Send + 'static,
M: SessionManager,
{
pub fn new(
service_factory: impl Fn() -> Result<S, std::io::Error> + Send + Sync + 'static,
session_manager: Arc<M>,
config: StreamableHttpServerConfig,
) -> Self {
Self {
config,
session_manager,
service_factory: Arc::new(service_factory),
}
}
fn get_service(&self) -> Result<S, std::io::Error> {
(self.service_factory)()
}
pub async fn handle<B>(&self, request: Request<B>) -> Response<BoxBody<Bytes, Infallible>>
where
B: Body + Send + 'static,
B::Error: Display,
{
let method = request.method().clone();
let allowed_methods = match self.config.stateful_mode {
true => "GET, POST, DELETE",
false => "POST",
};
let result = match (method, self.config.stateful_mode) {
(Method::POST, _) => self.handle_post(request).await,
// if we're not in stateful mode, we don't support GET or DELETE because there is no session
(Method::GET, true) => self.handle_get(request).await,
(Method::DELETE, true) => self.handle_delete(request).await,
_ => {
// Handle other methods or return an error
let response = Response::builder()
.status(http::StatusCode::METHOD_NOT_ALLOWED)
.header(ALLOW, allowed_methods)
.body(Full::new(Bytes::from("Method Not Allowed")).boxed())
.expect("valid response");
return response;
}
};
match result {
Ok(response) => response,
Err(response) => response,
}
}
async fn handle_get<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
where
B: Body + Send + 'static,
B::Error: Display,
{
// check accept header
if !request
.headers()
.get(http::header::ACCEPT)
.and_then(|header| header.to_str().ok())
.is_some_and(|header| header.contains(EVENT_STREAM_MIME_TYPE))
{
return Ok(Response::builder()
.status(http::StatusCode::NOT_ACCEPTABLE)
.body(
Full::new(Bytes::from(
"Not Acceptable: Client must accept text/event-stream",
))
.boxed(),
)
.expect("valid response"));
}
// check session id
let session_id = request
.headers()
.get(HEADER_SESSION_ID)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned().into());
let Some(session_id) = session_id else {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(Full::new(Bytes::from("Unauthorized: Session ID is required")).boxed())
.expect("valid response"));
};
// check if session exists
let has_session = self
.session_manager
.has_session(&session_id)
.await
.map_err(internal_error_response("check session"))?;
if !has_session {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(Full::new(Bytes::from("Unauthorized: Session not found")).boxed())
.expect("valid response"));
}
// check if last event id is provided
let last_event_id = request
.headers()
.get(HEADER_LAST_EVENT_ID)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned());
if let Some(last_event_id) = last_event_id {
// check if session has this event id
let stream = self
.session_manager
.resume(&session_id, last_event_id)
.await
.map_err(internal_error_response("resume session"))?;
// Resume doesn't need priming - client already has the event ID
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
} else {
// create standalone stream
let stream = self
.session_manager
.create_standalone_stream(&session_id)
.await
.map_err(internal_error_response("create standalone stream"))?;
// Prepend priming event if sse_retry configured
let stream = if let Some(retry) = self.config.sse_retry {
let priming = ServerSseMessage {
event_id: Some("0".into()),
message: None,
retry: Some(retry),
};
futures::stream::once(async move { priming })
.chain(stream)
.left_stream()
} else {
stream.right_stream()
};
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
}
}
async fn handle_post<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
where
B: Body + Send + 'static,
B::Error: Display,
{
// check accept header
if !request
.headers()
.get(http::header::ACCEPT)
.and_then(|header| header.to_str().ok())
.is_some_and(|header| {
header.contains(JSON_MIME_TYPE) && header.contains(EVENT_STREAM_MIME_TYPE)
})
{
return Ok(Response::builder()
.status(http::StatusCode::NOT_ACCEPTABLE)
.body(Full::new(Bytes::from("Not Acceptable: Client must accept both application/json and text/event-stream")).boxed())
.expect("valid response"));
}
// check content type
if !request
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|header| header.to_str().ok())
.is_some_and(|header| header.starts_with(JSON_MIME_TYPE))
{
return Ok(Response::builder()
.status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
.body(
Full::new(Bytes::from(
"Unsupported Media Type: Content-Type must be application/json",
))
.boxed(),
)
.expect("valid response"));
}
// json deserialize request body
let (part, body) = request.into_parts();
let mut message = match expect_json(body).await {
Ok(message) => message,
Err(response) => return Ok(response),
};
if self.config.stateful_mode {
// do we have a session id?
let session_id = part
.headers
.get(HEADER_SESSION_ID)
.and_then(|v| v.to_str().ok());
if let Some(session_id) = session_id {
let session_id = session_id.to_owned().into();
let has_session = self
.session_manager
.has_session(&session_id)
.await
.map_err(internal_error_response("check session"))?;
if !has_session {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(Full::new(Bytes::from("Unauthorized: Session not found")).boxed())
.expect("valid response"));
}
// inject request part to extensions
match &mut message {
ClientJsonRpcMessage::Request(req) => {
req.request.extensions_mut().insert(part);
}
ClientJsonRpcMessage::Notification(not) => {
not.notification.extensions_mut().insert(part);
}
_ => {
// skip
}
}
match message {
ClientJsonRpcMessage::Request(_) => {
let stream = self
.session_manager
.create_stream(&session_id, message)
.await
.map_err(internal_error_response("get session"))?;
// Prepend priming event if sse_retry configured
let stream = if let Some(retry) = self.config.sse_retry {
let priming = ServerSseMessage {
event_id: Some("0".into()),
message: None,
retry: Some(retry),
};
futures::stream::once(async move { priming })
.chain(stream)
.left_stream()
} else {
stream.right_stream()
};
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
}
ClientJsonRpcMessage::Notification(_)
| ClientJsonRpcMessage::Response(_)
| ClientJsonRpcMessage::Error(_) => {
// handle notification
self.session_manager
.accept_message(&session_id, message)
.await
.map_err(internal_error_response("accept message"))?;
Ok(accepted_response())
}
}
} else {
let (session_id, transport) = self
.session_manager
.create_session()
.await
.map_err(internal_error_response("create session"))?;
if let ClientJsonRpcMessage::Request(req) = &mut message {
if !matches!(req.request, ClientRequest::InitializeRequest(_)) {
return Err(unexpected_message_response("initialize request"));
}
// inject request part to extensions
req.request.extensions_mut().insert(part);
} else {
return Err(unexpected_message_response("initialize request"));
}
let service = self
.get_service()
.map_err(internal_error_response("get service"))?;
// spawn a task to serve the session
tokio::spawn({
let session_manager = self.session_manager.clone();
let session_id = session_id.clone();
async move {
let service = serve_server::<S, M::Transport, _, TransportAdapterIdentity>(
service, transport,
)
.await;
match service {
Ok(service) => {
// on service created
let _ = service.waiting().await;
}
Err(e) => {
tracing::error!("Failed to create service: {e}");
}
}
let _ = session_manager
.close_session(&session_id)
.await
.inspect_err(|e| {
tracing::error!("Failed to close session {session_id}: {e}");
});
}
});
// get initialize response
let response = self
.session_manager
.initialize_session(&session_id, message)
.await
.map_err(internal_error_response("create stream"))?;
let stream = futures::stream::once(async move {
ServerSseMessage {
event_id: None,
message: Some(Arc::new(response)),
retry: None,
}
});
// Prepend priming event if sse_retry configured
let stream = if let Some(retry) = self.config.sse_retry {
let priming = ServerSseMessage {
event_id: Some("0".into()),
message: None,
retry: Some(retry),
};
futures::stream::once(async move { priming })
.chain(stream)
.left_stream()
} else {
stream.right_stream()
};
let mut response = sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
);
response.headers_mut().insert(
HEADER_SESSION_ID,
session_id
.parse()
.map_err(internal_error_response("create session id header"))?,
);
Ok(response)
}
} else {
let service = self
.get_service()
.map_err(internal_error_response("get service"))?;
match message {
ClientJsonRpcMessage::Request(mut request) => {
request.request.extensions_mut().insert(part);
let (transport, receiver) =
OneshotTransport::<RoleServer>::new(ClientJsonRpcMessage::Request(request));
let service = serve_directly(service, transport, None);
tokio::spawn(async move {
// on service created
let _ = service.waiting().await;
});
// Stateless mode: no priming (no session to resume)
let stream = ReceiverStream::new(receiver).map(|message| {
tracing::info!(?message);
ServerSseMessage {
event_id: None,
message: Some(Arc::new(message)),
retry: None,
}
});
Ok(sse_stream_response(
stream,
self.config.sse_keep_alive,
self.config.cancellation_token.child_token(),
))
}
ClientJsonRpcMessage::Notification(_notification) => {
// ignore
Ok(accepted_response())
}
ClientJsonRpcMessage::Response(_json_rpc_response) => Ok(accepted_response()),
ClientJsonRpcMessage::Error(_json_rpc_error) => Ok(accepted_response()),
}
}
}
async fn handle_delete<B>(&self, request: Request<B>) -> Result<BoxResponse, BoxResponse>
where
B: Body + Send + 'static,
B::Error: Display,
{
// check session id
let session_id = request
.headers()
.get(HEADER_SESSION_ID)
.and_then(|v| v.to_str().ok())
.map(|s| s.to_owned().into());
let Some(session_id) = session_id else {
// unauthorized
return Ok(Response::builder()
.status(http::StatusCode::UNAUTHORIZED)
.body(Full::new(Bytes::from("Unauthorized: Session ID is required")).boxed())
.expect("valid response"));
};
// close session
self.session_manager
.close_session(&session_id)
.await
.map_err(internal_error_response("close session"))?;
Ok(accepted_response())
}
}
codex-rs/vendor/rmcp/src/transport/worker.rs
Vendored
+208
View File
@@ -0,0 +1,208 @@
use std::borrow::Cow;
use tokio_util::sync::CancellationToken;
use tracing::{Instrument, Level};
use super::{IntoTransport, Transport};
use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
#[derive(Debug, thiserror::Error)]
pub enum WorkerQuitReason<E> {
#[error("Join error {0}")]
Join(#[from] tokio::task::JoinError),
#[error("Transport fatal {error}, when {context}")]
Fatal {
error: E,
context: Cow<'static, str>,
},
#[error("Transport cancelled")]
Cancelled,
#[error("Transport closed")]
TransportClosed,
#[error("Handler terminated")]
HandlerTerminated,
}
impl<E: std::error::Error + Send + 'static> WorkerQuitReason<E> {
pub fn fatal(error: E, context: impl Into<Cow<'static, str>>) -> Self {
Self::Fatal {
error,
context: context.into(),
}
}
pub fn fatal_context(context: impl Into<Cow<'static, str>>) -> impl FnOnce(E) -> Self {
|e| Self::Fatal {
error: e,
context: context.into(),
}
}
}
pub trait Worker: Sized + Send + 'static {
type Error: std::error::Error + Send + Sync + 'static;
type Role: ServiceRole;
fn err_closed() -> Self::Error;
fn err_join(e: tokio::task::JoinError) -> Self::Error;
fn run(
self,
context: WorkerContext<Self>,
) -> impl Future<Output = Result<(), WorkerQuitReason<Self::Error>>> + Send;
fn config(&self) -> WorkerConfig {
WorkerConfig::default()
}
}
pub struct WorkerSendRequest<W: Worker> {
pub message: TxJsonRpcMessage<W::Role>,
pub responder: tokio::sync::oneshot::Sender<Result<(), W::Error>>,
}
pub struct WorkerTransport<W: Worker> {
rx: tokio::sync::mpsc::Receiver<RxJsonRpcMessage<W::Role>>,
send_service: tokio::sync::mpsc::Sender<WorkerSendRequest<W>>,
join_handle: Option<tokio::task::JoinHandle<Result<(), WorkerQuitReason<W::Error>>>>,
_drop_guard: tokio_util::sync::DropGuard,
ct: CancellationToken,
}
pub struct WorkerConfig {
pub name: Option<String>,
pub channel_buffer_capacity: usize,
}
impl Default for WorkerConfig {
fn default() -> Self {
Self {
name: None,
channel_buffer_capacity: 16,
}
}
}
pub enum WorkerAdapter {}
impl<W: Worker> IntoTransport<W::Role, W::Error, WorkerAdapter> for W {
fn into_transport(self) -> impl Transport<W::Role, Error = W::Error> + 'static {
WorkerTransport::spawn(self)
}
}
impl<W: Worker> WorkerTransport<W> {
pub fn cancel_token(&self) -> CancellationToken {
self.ct.clone()
}
pub fn spawn(worker: W) -> Self {
Self::spawn_with_ct(worker, CancellationToken::new())
}
pub fn spawn_with_ct(worker: W, transport_task_ct: CancellationToken) -> Self {
let config = worker.config();
let worker_name = config.name;
let (to_transport_tx, from_handler_rx) =
tokio::sync::mpsc::channel::<WorkerSendRequest<W>>(config.channel_buffer_capacity);
let (to_handler_tx, from_transport_rx) =
tokio::sync::mpsc::channel::<RxJsonRpcMessage<W::Role>>(config.channel_buffer_capacity);
let context = WorkerContext {
to_handler_tx,
from_handler_rx,
cancellation_token: transport_task_ct.clone(),
};
let join_handle = tokio::spawn(async move {
worker
.run(context)
.instrument(tracing::span!(
Level::TRACE,
"transport_worker",
name = worker_name,
))
.await
.inspect_err(|e| match e {
WorkerQuitReason::Cancelled
| WorkerQuitReason::TransportClosed
| WorkerQuitReason::HandlerTerminated => {
tracing::debug!("worker quit with reason: {:?}", e);
}
WorkerQuitReason::Join(e) => {
tracing::error!("worker quit with join error: {:?}", e);
}
WorkerQuitReason::Fatal { error, context } => {
tracing::error!("worker quit with fatal: {error}, when {context}");
}
})
.inspect(|_| {
tracing::debug!("worker quit");
})
});
Self {
rx: from_transport_rx,
send_service: to_transport_tx,
join_handle: Some(join_handle),
ct: transport_task_ct.clone(),
_drop_guard: transport_task_ct.drop_guard(),
}
}
}
pub struct SendRequest<W: Worker> {
pub message: TxJsonRpcMessage<W::Role>,
pub responder: tokio::sync::oneshot::Sender<RxJsonRpcMessage<W::Role>>,
}
pub struct WorkerContext<W: Worker> {
pub to_handler_tx: tokio::sync::mpsc::Sender<RxJsonRpcMessage<W::Role>>,
pub from_handler_rx: tokio::sync::mpsc::Receiver<WorkerSendRequest<W>>,
pub cancellation_token: CancellationToken,
}
impl<W: Worker> WorkerContext<W> {
pub async fn send_to_handler(
&mut self,
item: RxJsonRpcMessage<W::Role>,
) -> Result<(), WorkerQuitReason<W::Error>> {
self.to_handler_tx
.send(item)
.await
.map_err(|_| WorkerQuitReason::HandlerTerminated)
}
pub async fn recv_from_handler(
&mut self,
) -> Result<WorkerSendRequest<W>, WorkerQuitReason<W::Error>> {
self.from_handler_rx
.recv()
.await
.ok_or(WorkerQuitReason::HandlerTerminated)
}
}
impl<W: Worker> Transport<W::Role> for WorkerTransport<W> {
type Error = W::Error;
fn send(
&mut self,
item: TxJsonRpcMessage<W::Role>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
let tx = self.send_service.clone();
let (responder, receiver) = tokio::sync::oneshot::channel();
let request = WorkerSendRequest {
message: item,
responder,
};
async move {
tx.send(request).await.map_err(|_| W::err_closed())?;
receiver.await.map_err(|_| W::err_closed())??;
Ok(())
}
}
async fn receive(&mut self) -> Option<RxJsonRpcMessage<W::Role>> {
self.rx.recv().await
}
async fn close(&mut self) -> Result<(), Self::Error> {
if let Some(handle) = self.join_handle.take() {
self.ct.cancel();
let _quit_reason = handle.await.map_err(W::err_join)?;
Ok(())
} else {
Ok(())
}
}
}
codex-rs/vendor/rmcp/src/transport/ws.rs
Vendored
+1
View File
@@ -0,0 +1 @@
// Maybe we don't really need a ws implementation?