public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

first pass of merging in warp (doesn't build)

Browse Source
This commit is contained in:
Ryan Ward
2026-07-01 16:08:58 -05:00
parent 2f64909469
commit 4770ac06b5
3662 changed files with 414574 additions and 89772 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/mcp/src/lib.rs
+95
View File
@@ -0,0 +1,95 @@
#[cfg(not(target_family = "wasm"))]
pub mod oauth;
#[cfg(not(target_family = "wasm"))]
pub mod runtime;
#[cfg(not(target_family = "wasm"))]
pub mod sse_transport;
use uuid::Uuid;
/// Information about a single connected MCP server.
pub struct TemplatableMCPServerInfo {
name: String,
service: rmcp::service::RunningService<
rmcp::RoleClient,
Box<dyn rmcp::service::DynService<rmcp::RoleClient>>,
>,
resources: Vec<rmcp::model::Resource>,
tools: Vec<rmcp::model::Tool>,
installation_id: Uuid,
description: Option<String>,
/// Whether the underlying transport uses authentication.
///
/// TODO(vorporeal): Use this to display a toast when server authentication and connection is complete, and
/// to provide a "log out" button.
#[allow(dead_code)]
is_authenticated_transport: bool,
}
impl TemplatableMCPServerInfo {
pub fn name(&self) -> &str {
&self.name
}
pub fn resources(&self) -> &Vec<rmcp::model::Resource> {
&self.resources
}
pub fn tools(&self) -> &Vec<rmcp::model::Tool> {
&self.tools
}
pub fn installation_id(&self) -> Uuid {
self.installation_id
}
pub fn description(&self) -> Option<&str> {
self.description.as_deref()
}
pub fn peer(&self) -> rmcp::Peer<rmcp::RoleClient> {
self.service.clone()
}
pub fn peer_if_connected(&self) -> Option<rmcp::Peer<rmcp::RoleClient>> {
if self.service.is_transport_closed() {
None
} else {
Some(self.service.clone())
}
}
pub fn has_tool(&self, tool_name: &str) -> bool {
self.tools.iter().any(|tool| tool.name == tool_name)
}
pub fn has_resource(&self, resource: &rmcp::model::Resource) -> bool {
self.resources
.iter()
.any(|other_resource| resource.uri == other_resource.uri)
}
pub fn has_resource_name_or_uri(&self, name: &str, uri: Option<&str>) -> bool {
self.resources.iter().any(|resource| {
if let Some(uri) = uri {
resource.uri == uri
} else {
resource.name == name
}
})
}
pub fn tool_input_schema(
&self,
tool_name: &str,
) -> Option<std::sync::Arc<rmcp::model::JsonObject>> {
self.tools
.iter()
.find(|tool| tool.name == tool_name)
.map(|tool| tool.input_schema.clone())
}
pub async fn shutdown(self) -> Result<rmcp::service::QuitReason, tokio::task::JoinError> {
self.service.cancel().await
}
}
crates/mcp/src/oauth.rs
+402
View File
@@ -0,0 +1,402 @@
use std::collections::HashMap;
use futures::future::BoxFuture;
use oauth2::{RefreshToken, TokenResponse as _};
use rmcp::transport::auth::{
AuthClient, AuthorizationManager, CredentialStore, InMemoryCredentialStore, OAuthClientConfig,
OAuthState, StoredCredentials,
};
use rmcp::transport::{AuthError, AuthorizationSession};
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use url::Url;
use uuid::Uuid;
use warp_core::channel::ChannelState;
use warpui_extras::secure_storage::AppContextExt as _;
pub const TEMPLATABLE_MCP_CREDENTIALS_KEY: &str = "TemplatableMcpCredentials";
pub const FILE_BASED_MCP_CREDENTIALS_KEY: &str = "FileBasedMcpCredentials";
/// The issuer URL for GitHub's OAuth provider.
const GITHUB_ISSUER: &str = "https://github.com/login/oauth";
static GITHUB_OAUTH_SCOPES: [&str; 7] = [
"repo",
"read:org",
"gist",
"notifications",
"user",
"project",
"workflow",
];
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PersistedCredentials {
/// The credential information that `rmcp` wants us to store and retrieve.
#[serde(flatten)]
credentials: StoredCredentials,
/// The client secret for the OAuth application.
///
/// This is needed to properly refresh tokens when using DCR (Dynamic Client Registration),
/// as the server expects the client to provide the secret when refreshing.
client_secret: Option<String>,
}
/// Maps cloud MCP installation UUID to its OAuth credentials in secure storage.
pub type PersistedCredentialsMap = HashMap<Uuid, PersistedCredentials>;
// Maps a consistent hash of the installation to its persisted credentials
pub type FileBasedPersistedCredentialsMap = HashMap<u64, PersistedCredentials>;
pub type PersistCredentialsCallback =
Box<dyn Fn(Uuid, PersistedCredentials) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
pub type RequiresAuthenticationCallback =
Box<dyn Fn(Uuid, String, String) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
pub type AuthenticatedCallback =
Box<dyn Fn(String) -> BoxFuture<'static, anyhow::Result<()>> + Send>;
/// A credential store that wraps [`InMemoryCredentialStore`] and persists token
/// updates to Warp's secure storage via a channel.
///
/// When rmcp auto-refreshes an expired access token at runtime, the rotated
/// tokens are only saved to the in-memory store by default. This wrapper
/// ensures they also get written back to secure storage so they survive app
/// restarts.
struct PersistingCredentialStore {
inner: InMemoryCredentialStore,
client_secret: Option<String>,
persist_tx: async_channel::Sender<PersistedCredentials>,
}
impl PersistingCredentialStore {
/// Per RFC 6749 §6, the authorization server MAY issue a new refresh token on
/// refresh, but is not required to. Many OAuth providers (e.g. Figma) only
/// issue a refresh token on the initial authorization grant and omit it from
/// subsequent refresh responses. If we blindly persist the new token response,
/// the refresh token is lost and the next session (or next in-process refresh)
/// requires a full re-auth.
///
/// When the new response omits a refresh token, carry forward the one already
/// in the store. See: <https://datatracker.ietf.org/doc/html/rfc6749#section-6>
async fn apply_refresh_token_carry_forward(&self, credentials: &mut StoredCredentials) {
if credentials
.token_response
.as_ref()
.is_none_or(|tr| tr.refresh_token().is_some())
{
return;
}
if let Some(prev_rt) = self
.inner
.load()
.await
.ok()
.and_then(|opt| opt)
.and_then(|prev| prev.token_response)
.and_then(|prev_tr| prev_tr.refresh_token().cloned())
{
if let Some(tr) = credentials.token_response.as_mut() {
// Carry forward the existing/previous refresh token, constructing new if needed
tr.set_refresh_token(Some(RefreshToken::new(prev_rt.secret().to_string())));
}
}
}
}
#[async_trait::async_trait]
impl CredentialStore for PersistingCredentialStore {
async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
self.inner.load().await
}
async fn save(&self, mut credentials: StoredCredentials) -> Result<(), AuthError> {
self.apply_refresh_token_carry_forward(&mut credentials)
.await;
self.inner.save(credentials.clone()).await?;
// Only persist credentials if we actually have any.
if credentials.token_response.is_some() {
let _ = self.persist_tx.try_send(PersistedCredentials {
credentials,
client_secret: self.client_secret.clone(),
});
}
Ok(())
}
async fn clear(&self) -> Result<(), AuthError> {
self.inner.clear().await
}
}
/// Installs a [`PersistingCredentialStore`] on the given auth manager so that
/// runtime token auto-refreshes are written back to Warp's secure storage.
///
/// A background tokio task is spawned to receive credential updates and persist
/// them via the provided callback. The task terminates when the auth manager (and
/// thus the credential store's sender) is dropped.
///
/// Note: this store is not responsible for the initial population of credentials.
/// Instead, the caller seeds the inner store with any existing credentials prior
/// to installation (see [`install_persisting_credential_store`]). This store's
/// sole role is to write token updates back to secure storage as they occur.
async fn install_persisting_credential_store(
auth_manager: &mut AuthorizationManager,
persisted_credentials: Option<PersistedCredentials>,
persist_credentials: PersistCredentialsCallback,
installation_uuid: Uuid,
) {
let client_secret = persisted_credentials
.as_ref()
.and_then(|c| c.client_secret.clone());
let in_memory_store = InMemoryCredentialStore::new();
// If we have persisted credentials, populate the backing in-memory store with them.
if let Some(credentials) = persisted_credentials {
let _ = in_memory_store.save(credentials.credentials).await;
}
let (persist_tx, persist_rx) = async_channel::unbounded();
let store = PersistingCredentialStore {
inner: in_memory_store,
client_secret,
persist_tx,
};
auth_manager.set_credential_store(store);
tokio::spawn(async move {
while let Ok(credentials) = persist_rx.recv().await {
if let Err(err) = persist_credentials(installation_uuid, credentials).await {
log::warn!("Failed to persist auto-refreshed MCP credentials: {err:?}");
}
}
});
}
/// Context for OAuth authentication flows.
pub struct AuthContext {
pub oauth_result_rx: async_channel::Receiver<CallbackResult>,
pub uuid: Uuid,
pub persisted_credentials: Option<PersistedCredentials>,
/// Whether the client is running in headless/CLI mode.
pub is_headless: bool,
/// Whether this server was auto-discovered from a repo MCP configuration file.
pub is_file_based: bool,
pub persist_credentials: PersistCredentialsCallback,
pub requires_authentication: RequiresAuthenticationCallback,
pub authenticated: Option<AuthenticatedCallback>,
}
/// Result of OAuth callback.
#[derive(Debug, Clone)]
pub enum CallbackResult {
Success { code: String, csrf_token: String },
Error { error: Option<String> },
}
/// Makes an authenticated client for the given authorization server.
///
/// This takes in the URL of the resource to authenticate for, and uses that
/// to determine the authorization server.
///
/// Upon success, returns the client and a boolean indicating whether the user was required to
/// re-authenticate (e.g. re-log in).
pub async fn make_authenticated_client(
resource_url: &str,
auth_context: AuthContext,
) -> Result<(AuthClient<reqwest::Client>, bool), AuthError> {
let AuthContext {
oauth_result_rx,
uuid,
persisted_credentials,
is_headless,
is_file_based,
persist_credentials,
requires_authentication,
..
} = auth_context;
// Build the redirect URI using the channel's URL scheme.
// Routing data (the server UUID) is passed via the OAuth `state` parameter instead
// of the redirect URI so that the URI exactly matches what is registered during
// Dynamic Client Registration, satisfying RFC 6749 §3.1.2.2 exact-match validation.
let redirect_uri = format!("{}://mcp/oauth2callback", ChannelState::url_scheme());
// Create the auth manager and initialize it with a backing credential store that persists
// new credentials to secure storage.
let client_id = persisted_credentials
.as_ref()
.map(|c| c.credentials.client_id.clone());
let client_secret = persisted_credentials
.as_ref()
.and_then(|c| c.client_secret.clone());
let mut auth_manager = AuthorizationManager::new(resource_url).await?;
install_persisting_credential_store(
&mut auth_manager,
persisted_credentials,
persist_credentials,
uuid,
)
.await;
// If we loaded persisted credentials from the store, and we have a valid access token
// (or successfully refreshed a valid refresh token), we're already authorized and good
// to go.
if auth_manager.initialize_from_store().await? && auth_manager.get_access_token().await.is_ok()
{
if let (Some(client_id), Some(client_secret)) = (client_id, client_secret) {
auth_manager.configure_client(
OAuthClientConfig::new(client_id, redirect_uri.clone())
.with_client_secret(client_secret),
)?;
}
return Ok((AuthClient::new(reqwest::Client::new(), auth_manager), false));
}
// If we're in headless mode and we reach here, it means we either have no credentials
// or the cached credentials failed to refresh. Block interactive OAuth in headless mode.
if is_headless {
if is_file_based {
log::warn!(
"File-based MCP server {uuid} requires OAuth authentication; \
skipping in headless mode. To use this server, authenticate it \
in the Warp desktop app first."
);
}
return Err(AuthError::AuthorizationFailed(
"MCP server requires OAuth authentication. Please authenticate this server in the \
Warp desktop app first, then try again."
.to_string(),
));
}
let metadata = auth_manager.discover_metadata().await?;
// Configure the auth manager's OAuth client using dynamic or static client registration.
let mut oauth_state = if let Some(provider) = metadata
.issuer
.as_deref()
.and_then(ChannelState::mcp_oauth_provider_by_issuer)
{
// Configure the auth manager based on the static MCP configuration for this
// issuer.
auth_manager.set_metadata(metadata);
let scopes = if provider.issuer == GITHUB_ISSUER {
GITHUB_OAUTH_SCOPES
.into_iter()
.map(ToString::to_string)
.collect()
} else {
vec![]
};
auth_manager.configure_client(
OAuthClientConfig::new(provider.client_id, redirect_uri.clone())
.with_client_secret(provider.client_secret)
.with_scopes(scopes),
)?;
// We do a scope "upgrade" with no additional scopes here as it's the easiest way
// to construct an auth URL.
let auth_url = auth_manager.request_scope_upgrade("").await?;
OAuthState::Session(AuthorizationSession::for_scope_upgrade(
auth_manager,
auth_url,
&redirect_uri,
))
} else {
// Try dynamic client registration.
let mut oauth_state = OAuthState::Unauthorized(auth_manager);
oauth_state
.start_authorization(&[], &redirect_uri, Some("Warp"))
.await?;
oauth_state
};
let auth_url = oauth_state.get_authorization_url().await?;
// Extract the CSRF token that rmcp embedded as the `state` query parameter in the
// authorization URL. We register a csrf→uuid mapping on the manager so that
// `handle_oauth_callback` can route the callback to the right server without
// relying on `server_id` being present in the redirect URI.
let csrf_state = Url::parse(&auth_url)
.ok()
.and_then(|u| {
u.query_pairs()
.find(|(k, _)| k == "state")
.map(|(_, v)| v.into_owned())
})
.unwrap_or_default();
if let Err(err) = requires_authentication(uuid, csrf_state, auth_url).await {
log::warn!("Failed to emit RequiresAuthentication state: {err:?}");
}
// Wait for the authorization code from the OAuth callback channel.
let oauth_result = oauth_result_rx
.recv()
.await
.map_err(|e| AuthError::InternalError(e.to_string()))?;
let (code, csrf_token) = match &oauth_result {
CallbackResult::Success { code, csrf_token } => (code, csrf_token),
CallbackResult::Error { error } => {
return Err(AuthError::AuthorizationFailed(
error.as_deref().unwrap_or("unknown error").to_string(),
));
}
};
// Handle the callback with the received authorization code and CSRF token.
oauth_state.handle_callback(code, csrf_token).await?;
let auth_manager = oauth_state.into_authorization_manager().ok_or_else(|| {
AuthError::InternalError("Failed to create authorization manager".to_string())
})?;
Ok((AuthClient::new(reqwest::Client::new(), auth_manager), true))
}
/// Loads credentials from secure storage at the provided key.
pub fn load_credentials_from_secure_storage<T: DeserializeOwned + Default>(
app: &mut warpui::AppContext,
key: &str,
) -> T {
app.secure_storage()
.read_value(key)
.inspect_err(|err| {
if !matches!(err, warpui_extras::secure_storage::Error::NotFound) {
log::warn!("Failed to read MCP credentials from secure storage: {err:#}");
}
})
.ok()
.and_then(|value| serde_json::from_str(&value).ok())
.unwrap_or_default()
}
/// Writes credentials to secure storage at the provided key.
pub fn write_to_secure_storage<T: Serialize>(
app: &mut warpui::AppContext,
key: &str,
credentials: &T,
) {
match serde_json::to_string(credentials) {
Ok(json) => {
app.secure_storage()
.write_value(key, &json)
.inspect_err(|err| {
log::error!("Failed to write MCP credentials to secure storage: {err:#}")
})
.ok();
}
Err(err) => {
log::error!("Failed to serialize MCP credentials for secure storage: {err:#}");
}
}
}
#[cfg(test)]
#[path = "oauth_tests.rs"]
mod tests;
crates/mcp/src/oauth_tests.rs
+172
View File
@@ -0,0 +1,172 @@
use rmcp::transport::auth::OAuthTokenResponse;
use super::*;
/// Builds a minimal `OAuthTokenResponse` for tests, optionally with a refresh token.
fn make_test_token_response(refresh_token: Option<&str>) -> OAuthTokenResponse {
let mut json = serde_json::json!({
"access_token": "test_access_token",
"token_type": "bearer",
"expires_in": 3600,
});
if let Some(rt) = refresh_token {
json["refresh_token"] = serde_json::Value::String(rt.to_string());
}
serde_json::from_value(json).expect("OAuthTokenResponse deserialization")
}
/// Constructs a fresh `PersistingCredentialStore` plus the receiver side of its
/// persist channel so tests can observe what would be written to secure storage.
fn make_test_store(
client_secret: Option<String>,
) -> (
PersistingCredentialStore,
async_channel::Receiver<PersistedCredentials>,
) {
let (tx, rx) = async_channel::unbounded();
let store = PersistingCredentialStore {
inner: InMemoryCredentialStore::new(),
client_secret,
persist_tx: tx,
};
(store, rx)
}
/// Backward compatibility: credentials persisted by older Warp versions do not
/// have the `token_received_at` field. Deserializing them must succeed and
/// default to `None` so the next refresh can populate it. Failing this test
/// would mean every existing user loses their MCP OAuth tokens on upgrade.
#[test]
fn persisted_credentials_deserializes_legacy_format_without_received_at() {
// Legacy format: no `token_received_at` field.
let legacy_json = r#"{
"client_id": "client-abc",
"client_secret": null,
"token_response": {
"access_token": "old_access",
"token_type": "bearer",
"expires_in": 3600,
"refresh_token": "old_refresh"
}
}"#;
let parsed: PersistedCredentials =
serde_json::from_str(legacy_json).expect("legacy format must deserialize");
assert_eq!(parsed.credentials.client_id, "client-abc");
assert_eq!(parsed.credentials.token_received_at, None);
}
/// Regression test for #8863. When rmcp persists refreshed credentials via
/// `CredentialStore::save`, the `token_received_at` must be forwarded into
/// the channel so the persisted (secure-storage) representation can stamp
/// it. Without this, a restart would lose the timestamp and rmcp's
/// pre-emptive refresh check would be permanently disabled for the cached
/// session.
#[tokio::test]
async fn save_forwards_token_received_at_to_persist_channel() {
let (store, rx) = make_test_store(Some("client_secret_xyz".to_string()));
let credentials = StoredCredentials::new(
"client-id".to_string(),
Some(make_test_token_response(Some("refresh-1"))),
Vec::new(),
Some(1_700_000_500),
);
store.save(credentials).await.expect("save succeeds");
let persisted = rx.try_recv().expect("persist channel received credentials");
assert_eq!(persisted.credentials.token_received_at, Some(1_700_000_500));
assert_eq!(persisted.credentials.client_id, "client-id");
assert_eq!(
persisted.client_secret.as_deref(),
Some("client_secret_xyz")
);
}
/// Defensive: if rmcp ever calls `save` without a `token_received_at`
/// (e.g., during initial credential set-up before refresh), we must
/// propagate `None` rather than silently substituting a value.
#[tokio::test]
async fn save_forwards_none_when_received_at_is_none() {
let (store, rx) = make_test_store(None);
let credentials = StoredCredentials::new(
"c".to_string(),
Some(make_test_token_response(None)),
Vec::new(),
None,
);
store.save(credentials).await.expect("save succeeds");
let persisted = rx.try_recv().expect("persist channel received credentials");
assert_eq!(persisted.credentials.token_received_at, None);
}
/// `save` only forwards a credentials snapshot to the persist channel when
/// `token_response` is `Some`. This guards the existing branch from regression.
#[tokio::test]
async fn save_skips_persist_when_token_response_absent() {
let (store, rx) = make_test_store(None);
let credentials =
StoredCredentials::new("c".to_string(), None, Vec::new(), Some(1_700_000_500));
store.save(credentials).await.expect("save succeeds");
assert!(
rx.try_recv().is_err(),
"no PersistedCredentials should be sent when token_response is absent"
);
}
/// The carry-forward of refresh tokens (when the OAuth server omits one
/// from a refresh response) must not interfere with `token_received_at`
/// propagation. Tests both behaviors in one save: the new credentials get
/// the prior refresh token AND the new `token_received_at`.
#[tokio::test]
async fn save_carries_forward_refresh_token_and_preserves_received_at() {
let (store, rx) = make_test_store(None);
// Seed the inner store with prior credentials that have a refresh token.
store
.inner
.save(StoredCredentials::new(
"c".to_string(),
Some(make_test_token_response(Some("prior-refresh-token"))),
Vec::new(),
Some(1_699_000_000),
))
.await
.expect("seed succeeds");
// Now save NEW credentials that omit a refresh token, simulating a
// refresh response from a server that does not rotate refresh tokens.
let new_credentials = StoredCredentials::new(
"c".to_string(),
Some(make_test_token_response(None)),
Vec::new(),
Some(1_700_000_500),
);
store.save(new_credentials).await.expect("save succeeds");
let persisted = rx.try_recv().expect("persist channel received credentials");
assert_eq!(
persisted.credentials.token_received_at,
Some(1_700_000_500),
"newer received_at preserved"
);
let refresh_token = persisted
.credentials
.token_response
.and_then(|tr| tr.refresh_token().cloned());
assert_eq!(
refresh_token.map(|rt| rt.secret().to_string()),
Some("prior-refresh-token".to_string()),
"prior refresh token carried forward"
);
}
crates/mcp/src/runtime.rs
+585
View File
@@ -0,0 +1,585 @@
//! Capability-gating helpers used during MCP server startup.
//!
//! Each `query_*_for` function pairs a capability check with the actual list
//! call from rmcp, gating the call on advertisement and failing soft on errors.
//! They take the list call as a closure so unit tests can drive the gate-and-
//! fail-soft control flow with a fake `RunningService` substitute.
use std::collections::HashMap;
use std::future::Future;
use cfg_if::cfg_if;
use cloud_object_models::{StaticEnvVar, TransportType};
use futures::FutureExt as _;
use rmcp::transport::ConfigureCommandExt as _;
use rmcp::ServiceExt as _;
use simple_logger::SimpleLogger;
use tokio::io::AsyncBufReadExt as _;
use uuid::Uuid;
use super::TemplatableMCPServerInfo;
type ReqwestHttpTransport = rmcp::transport::StreamableHttpClientTransport<reqwest::Client>;
type ReqwestSseTransport = crate::sse_transport::SseClientTransport<reqwest::Client>;
/// Convert an rmcp error to a user-friendly error message.
pub fn error_to_user_message(error: &rmcp::RmcpError) -> String {
match error {
rmcp::RmcpError::ClientInitialize(err) => {
format!("Failed to initialize client: {}", err)
}
rmcp::RmcpError::ServerInitialize(err) => {
format!("Failed to initialize server: {}", err)
}
rmcp::RmcpError::TransportCreation { error, .. } => {
format!("Failed to establish connection: {}", error)
}
rmcp::RmcpError::Runtime(err) => {
format!("Runtime error: {}", err)
}
rmcp::RmcpError::Service(err) => match err {
rmcp::ServiceError::McpError(_) => {
"Server returned an error. Please check server logs for details.".to_string()
}
rmcp::ServiceError::TransportSend(_) => {
"Failed to send data to server. Connection may have been lost.".to_string()
}
rmcp::ServiceError::TransportClosed => {
"Connection closed unexpectedly. The server may have crashed.".to_string()
}
rmcp::ServiceError::UnexpectedResponse => {
"Server sent an unexpected response. The server may be incompatible.".to_string()
}
rmcp::ServiceError::Cancelled { reason } => format!(
"Operation was cancelled with reason: {}",
reason.clone().unwrap_or("Unknown reason".to_string())
),
rmcp::ServiceError::Timeout { timeout } => {
format!(
"Connection timed out after {} seconds. The server may be unresponsive.",
timeout.as_secs()
)
}
_ => format!("Service error: {}", err),
},
// The enum is marked as non-exhaustive, so we need a catch-all.
_ => {
format!("Error: {error}")
}
}
}
/// Builds a `HeaderMap` from a `HashMap<String, String>` of user-provided headers.
///
/// Invalid header names or values are skipped.
fn build_header_map(headers: &HashMap<String, String>) -> reqwest::header::HeaderMap {
headers.try_into().unwrap_or_default()
}
/// Builds a reqwest client with custom headers for MCP HTTP/SSE connections.
#[allow(clippy::result_large_err)]
pub fn build_client_with_headers(
headers: &HashMap<String, String>,
) -> Result<reqwest::Client, rmcp::RmcpError> {
let header_map = build_header_map(headers);
reqwest::Client::builder()
.default_headers(header_map)
.build()
.map_err(|e| {
rmcp::RmcpError::transport_creation::<ReqwestHttpTransport>(format!(
"Failed to build client with headers: {e}",
))
})
}
/// Spawns a new MCP server from a given [`TransportType`].
#[allow(clippy::result_large_err)]
pub async fn spawn_server(
server_name: String,
description: Option<String>,
uuid: Uuid,
transport_type: TransportType,
logger: SimpleLogger,
auth_context: Option<crate::oauth::AuthContext>,
) -> Result<TemplatableMCPServerInfo, rmcp::RmcpError> {
logger.log("[note] Attention! There may be sensitive information (such as API keys) in these logs. Make sure to redact any secrets before sharing with others.".to_string());
let mut is_authenticated_transport = false;
let service = match transport_type {
TransportType::CLIServer(cli_server) => {
logger.log("[info] MCP: Using stdio transport".to_string());
cfg_if! {
if #[cfg(windows)] {
// We wrap the command in cmd.exe /c to allow Windows to be responsible for resolving the
// PATH variable rather than depending on the `Command` implementation, which only looks for
// `.exe` files in directories found in PATH.
// https://github.com/rust-lang/rust/issues/37519
let command = "cmd.exe".to_owned();
let args = std::iter::once("/c".to_owned())
.chain(std::iter::once(cli_server.command))
.chain(cli_server.args)
.collect::<Vec<String>>();
} else {
let command = cli_server.command;
let args = cli_server.args;
}
}
// Capture the command and configured cwd for diagnostics before they're
// moved into the Command builder closure.
let command_for_log = command.clone();
let cwd_for_log = cli_server.cwd_parameter.clone();
// Try to spawn the child process.
let (transport, stderr) = rmcp::transport::TokioChildProcess::builder(
tokio::process::Command::new(command).configure(|cmd| {
cmd.args(args);
if let Some(cwd) = cli_server.cwd_parameter {
cmd.current_dir(cwd);
}
for StaticEnvVar { name, value } in cli_server.static_env_vars.iter() {
if value.is_empty() {
// Skip empty/unset environment variables so that, in the CLI, they can be inherited.
logger.log(format!(
"[warn] MCP: Skipping empty environment variable: {name}"
));
continue;
}
cmd.env(name, value);
}
// On Windows, ensure that no console window is shown.
#[cfg(windows)]
cmd.creation_flags(windows::Win32::System::Threading::CREATE_NO_WINDOW.0);
}),
)
.stderr(std::process::Stdio::piped())
.spawn()
.map_err(|err| {
if err.kind() == std::io::ErrorKind::NotFound {
let cwd_display = cwd_for_log
.as_deref()
.unwrap_or("<inherited from Warp's process cwd>");
logger.log(format!(
"[error] MCP: Failed to spawn '{server_name}': command '{command_for_log}' \
not found (cwd: {cwd_display}). If your MCP server depends on a specific \
working directory, set the `working_directory` field in your config to \
override the default."
));
}
rmcp::RmcpError::transport_creation::<rmcp::transport::TokioChildProcess>(err)
})?;
let pid = transport
.id()
.map(|pid| pid.to_string())
.unwrap_or("??".to_string());
// We always expect to have an stderr, but this is marginally safer than unwrapping.
if let Some(stderr) = stderr {
let logger = logger.clone();
// Spawn a background task to forward from the child process's stderr to our logger.
tokio::spawn(async move {
let mut buf = String::new();
let mut reader = tokio::io::BufReader::new(stderr);
loop {
match reader.read_line(&mut buf).await {
// EOF.
Ok(0) => return,
// Read some data.
Ok(_) => logger.log(format!("[info] MCP [pid: {pid}] stderr: {buf}")),
// Failed to read from the child process's stderr.
Err(e) => {
log::error!("Failed to read stderr: {e}");
return;
}
}
}
});
}
// Wrap the transport in a logging wrapper.
let transport = TransportLoggingWrapper {
transport,
logger: logger.clone(),
};
// Create the MCP client and connect to the server.
Ok::<_, rmcp::RmcpError>(make_client_info().into_dyn().serve(transport).await?)
}
TransportType::ServerSentEvents(sse_server) => {
let headers: HashMap<String, String> = sse_server
.headers
.iter()
.map(|h| (h.name.clone(), h.value.clone()))
.collect();
match determine_transport(server_name.clone(), &sse_server.url, &headers, auth_context)
.await
{
// TODO: these need headers also?
Ok(Transport::Http(Some(client))) => {
is_authenticated_transport = true;
logger.log("[info] MCP: Using Streaming HTTP transport".to_string());
let transport = rmcp::transport::StreamableHttpClientTransport::with_client(
client,
rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig::with_uri(
sse_server.url.clone(),
),
);
let transport = TransportLoggingWrapper {
transport,
logger: logger.clone(),
};
Ok(make_client_info().into_dyn().serve(transport).await?)
}
Ok(Transport::Http(None)) => {
logger.log("[info] MCP: Using Streaming HTTP transport".to_string());
let transport = if headers.is_empty() {
rmcp::transport::StreamableHttpClientTransport::from_uri(
sse_server.url.clone(),
)
} else {
let client = build_client_with_headers(&headers)?;
rmcp::transport::StreamableHttpClientTransport::with_client(
client,
rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig::with_uri(
sse_server.url.clone(),
),
)
};
let transport = TransportLoggingWrapper {
transport,
logger: logger.clone(),
};
Ok(make_client_info().into_dyn().serve(transport).await?)
}
Ok(Transport::Sse(Some(client))) => {
is_authenticated_transport = true;
logger.log("[info] MCP: Using (legacy) SSE transport (due to preflight failing with a 404)".to_string());
let transport = crate::sse_transport::SseClientTransport::start_with_client(
client,
crate::sse_transport::SseClientConfig {
sse_endpoint: sse_server.url.into(),
..Default::default()
},
)
.await
.map_err(rmcp::RmcpError::transport_creation::<ReqwestSseTransport>)?;
let transport = TransportLoggingWrapper {
transport,
logger: logger.clone(),
};
Ok(make_client_info().into_dyn().serve(transport).await?)
}
Ok(Transport::Sse(None)) => {
logger.log("[info] MCP: Using (legacy) SSE transport (due to preflight failing with a 404)".to_string());
let transport = if headers.is_empty() {
crate::sse_transport::SseClientTransport::start(sse_server.url.clone())
.await
.map_err(|e| {
rmcp::RmcpError::transport_creation::<ReqwestSseTransport>(e)
})?
} else {
let client = build_client_with_headers(&headers)?;
crate::sse_transport::SseClientTransport::start_with_client(
client,
crate::sse_transport::SseClientConfig {
sse_endpoint: sse_server.url.clone().into(),
..Default::default()
},
)
.await
.map_err(rmcp::RmcpError::transport_creation::<ReqwestSseTransport>)?
};
let transport = TransportLoggingWrapper {
transport,
logger: logger.clone(),
};
Ok(make_client_info().into_dyn().serve(transport).await?)
}
Err(err) => {
logger.log(format!(
"[error] MCP: preflight connection to MCP server failed: {err:#}"
));
Err(err)?
}
}
}
}?;
let server_info = service.peer_info();
logger.log(format!("[info] MCP: Connected to server: {server_info:#?}"));
let capabilities = server_info.map(|info| &info.capabilities);
let resources =
query_resources_for(capabilities, &server_name, || service.list_all_resources()).await;
let tools = query_tools_for(capabilities, &server_name, || service.list_all_tools()).await;
Ok(TemplatableMCPServerInfo {
name: server_name,
service,
resources,
tools,
installation_id: uuid,
description,
is_authenticated_transport,
})
}
/// The transport to use for MCP.
enum Transport {
/// The HTTP transport, with an optional authenticated client.
Http(Option<rmcp::transport::auth::AuthClient<reqwest::Client>>),
/// The SSE transport, with an optional authenticated client.
Sse(Option<rmcp::transport::auth::AuthClient<reqwest::Client>>),
}
/// Determines which transport to use.
///
/// This sends a "preflight" InitializeRequest to the server to determine whether the
/// server supports the HTTP transport (or needs to use the SSE transport), and if
/// authentication is required.
#[allow(clippy::result_large_err)]
async fn determine_transport(
server_name: String,
url: &str,
headers: &HashMap<String, String>,
auth_context: Option<crate::oauth::AuthContext>,
) -> Result<Transport, rmcp::RmcpError> {
use reqwest::StatusCode;
fn unexpected_error(status: reqwest::StatusCode) -> rmcp::RmcpError {
rmcp::RmcpError::transport_creation::<ReqwestHttpTransport>(format!(
"Unexpected status code: {status}"
))
}
match send_initialize_request(url, headers, None).await? {
StatusCode::OK => Ok(Transport::Http(None)),
StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED => Ok(Transport::Sse(None)),
StatusCode::UNAUTHORIZED => {
let Some(mut auth_context) = auth_context else {
return Err(rmcp::RmcpError::transport_creation::<ReqwestHttpTransport>(
"Server requires authentication, which is not yet supported.".to_string(),
));
};
// Grab the post-authentication callback so we can invoke it once we know for sure that we successfully
// went through the OAuth flow for a server and were able to successfully send an initialize request.
let authenticated_callback = std::mem::take(&mut auth_context.authenticated);
// Go through the OAuth flow to get an authenticated client.
// This will first attempt to use cached credentials before starting interactive OAuth.
let (client, did_require_login) =
crate::oauth::make_authenticated_client(url, auth_context)
.await
.map_err(rmcp::RmcpError::transport_creation::<ReqwestHttpTransport>)?;
// Define a helper function to invoke when we've successfully authenticated.
let emit_authenticated_notification = async move || {
if did_require_login {
if let Some(authenticated_callback) = authenticated_callback {
if let Err(err) = authenticated_callback(server_name).await {
log::warn!("Failed to emit MCP authenticated notification: {err:?}");
}
}
}
};
match send_initialize_request(url, headers, Some(&client)).await? {
StatusCode::OK => {
emit_authenticated_notification().await;
Ok(Transport::Http(Some(client)))
}
StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED => {
emit_authenticated_notification().await;
Ok(Transport::Sse(Some(client)))
}
other => Err(unexpected_error(other)),
}
}
status => Err(unexpected_error(status)),
}
}
/// Sends an InitializeRequest to the server, and returns the HTTP status code from the response.
#[allow(clippy::result_large_err)]
async fn send_initialize_request(
url: &str,
headers: &HashMap<String, String>,
auth_client: Option<&rmcp::transport::auth::AuthClient<reqwest::Client>>,
) -> Result<reqwest::StatusCode, rmcp::RmcpError> {
use rmcp::transport::common::http_header::{EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE};
let request = rmcp::model::InitializeRequest::new(make_client_info());
let request = rmcp::model::ClientJsonRpcMessage::request(
rmcp::model::ClientRequest::InitializeRequest(request),
rmcp::model::RequestId::Number(0),
);
let mut request = build_client_with_headers(headers)?
.post(url)
.header(
http::header::ACCEPT,
[EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "),
)
.json(&request);
if let Some(auth_client) = auth_client.as_ref() {
let access_token = auth_client
.get_access_token()
.await
.map_err(rmcp::RmcpError::transport_creation::<ReqwestHttpTransport>)?;
request = request.bearer_auth(access_token);
}
let response = request
.send()
.await
.map_err(rmcp::RmcpError::transport_creation::<ReqwestHttpTransport>)?;
Ok(response.status())
}
/// Creates a [`ClientInfo`] for the MCP client.
///
/// This tells the MCP server who we are and what capabilities we have.
fn make_client_info() -> rmcp::model::ClientInfo {
rmcp::model::ClientInfo::new(
Default::default(),
rmcp::model::Implementation::new(
warp_core::channel::ChannelState::app_id().to_string(),
warp_core::channel::ChannelState::app_version()
.map(|v| v.to_string())
.unwrap_or_default(),
),
)
}
/// Whether to query `resources/list` for a server with the given capabilities.
///
/// Per the MCP spec, the client should only invoke a list method when the server
/// has advertised the corresponding capability during initialization.
fn should_query_resources(capabilities: Option<&rmcp::model::ServerCapabilities>) -> bool {
capabilities.is_some_and(|c| c.resources.is_some())
}
/// Whether to query `tools/list` for a server with the given capabilities.
///
/// Per the MCP spec, the client should only invoke a list method when the server
/// has advertised the corresponding capability during initialization.
fn should_query_tools(capabilities: Option<&rmcp::model::ServerCapabilities>) -> bool {
capabilities.is_some_and(|c| c.tools.is_some())
}
/// Query `resources/list` for a connected MCP server.
///
/// Skips the call entirely when `resources` was not advertised. Treats any
/// listing error as "no resources" (fail-soft) so a flaky `resources/list`
/// does not abort the entire server startup. Mirrors the behavior of
/// [`query_tools_for`] so the two capabilities are handled symmetrically.
async fn query_resources_for<F, Fut>(
capabilities: Option<&rmcp::model::ServerCapabilities>,
server_name: &str,
list_resources: F,
) -> Vec<rmcp::model::Resource>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<Vec<rmcp::model::Resource>, rmcp::ServiceError>>,
{
if !should_query_resources(capabilities) {
return Vec::new();
}
match list_resources().await {
Ok(result) => result,
Err(err) => {
log::warn!("Failed to list resources for MCP server '{server_name}': {err}");
Vec::new()
}
}
}
/// Query `tools/list` for a connected MCP server.
///
/// Skips the call entirely when `tools` was not advertised. Treats any listing
/// error as "no tools" (fail-soft) so a transient `tools/list` failure does
/// not abort the entire server startup — the user-visible regression #6798
/// was rooted in the prior asymmetric handling, where a tools-list error on
/// a server with healthy resources would propagate and fail startup.
async fn query_tools_for<F, Fut>(
capabilities: Option<&rmcp::model::ServerCapabilities>,
server_name: &str,
list_tools: F,
) -> Vec<rmcp::model::Tool>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<Vec<rmcp::model::Tool>, rmcp::ServiceError>>,
{
if !should_query_tools(capabilities) {
return Vec::new();
}
match list_tools().await {
Ok(result) => result,
Err(err) => {
log::warn!("Failed to list tools for MCP server '{server_name}': {err}");
Vec::new()
}
}
}
/// A wrapper around a [`rmcp::transport::Transport`] that logs all requests and responses.
struct TransportLoggingWrapper<T> {
transport: T,
logger: SimpleLogger,
}
impl<T: rmcp::transport::Transport<R>, R: rmcp::service::ServiceRole> rmcp::transport::Transport<R>
for TransportLoggingWrapper<T>
{
type Error = T::Error;
fn send(
&mut self,
item: rmcp::service::TxJsonRpcMessage<R>,
) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
if let Ok(json) = serde_json::to_string(&item) {
self.logger
.log(format!("[info] MCP: Sending request: {json}"));
}
let logger = self.logger.clone();
self.transport.send(item).map(move |result| {
if let Err(e) = &result {
logger.log(format!("[warn] MCP: Failed to send request: {e:#}"));
}
result
})
}
fn receive(
&mut self,
) -> impl Future<Output = Option<rmcp::service::RxJsonRpcMessage<R>>> + Send {
let logger = self.logger.clone();
async move {
let result = self.transport.receive().await;
if let Some(item) = &result {
if let Ok(json) = serde_json::to_string(item) {
logger.log(format!("[info] MCP: Received response: {json}"));
}
}
result
}
}
fn close(&mut self) -> impl Future<Output = Result<(), Self::Error>> + Send {
self.transport.close()
}
}
#[cfg(test)]
#[path = "runtime_tests.rs"]
mod tests;
crates/mcp/src/runtime_tests.rs
+281
View File
@@ -0,0 +1,281 @@
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use rmcp::model::{ErrorCode, ErrorData, Resource, ServerCapabilities, Tool};
use super::{query_resources_for, query_tools_for, should_query_resources, should_query_tools};
/// Build a `ServerCapabilities` with selected capability flags toggled on.
/// Each `Some(default)` mirrors how rmcp deserializes a capability the
/// server advertised with no inner flags set.
fn caps(tools: bool, resources: bool) -> ServerCapabilities {
match (tools, resources) {
(true, true) => ServerCapabilities::builder()
.enable_tools()
.enable_resources()
.build(),
(true, false) => ServerCapabilities::builder().enable_tools().build(),
(false, true) => ServerCapabilities::builder().enable_resources().build(),
(false, false) => ServerCapabilities::builder().build(),
}
}
fn test_tool(name: &str) -> Tool {
serde_json::from_value(serde_json::json!({
"name": name,
"description": "test tool",
"inputSchema": { "type": "object" },
}))
.expect("Tool deserialization")
}
fn test_resource(uri: &str) -> Resource {
serde_json::from_value(serde_json::json!({
"uri": uri,
"name": "test resource",
}))
.expect("Resource deserialization")
}
/// Regression test for warpdotdev/warp#6798: each capability is queried
/// independently. Previously, asymmetric handling could cause `tools/list`
/// to be skipped when a server advertised both `tools` and `resources`,
/// resulting in "No tools available" even though the server had tools.
#[test]
fn each_capability_is_queried_independently() {
for has_tools in [false, true] {
for has_resources in [false, true] {
let c = caps(has_tools, has_resources);
assert_eq!(
should_query_tools(Some(&c)),
has_tools,
"tools={has_tools}, resources={has_resources}",
);
assert_eq!(
should_query_resources(Some(&c)),
has_resources,
"tools={has_tools}, resources={has_resources}",
);
}
}
assert!(!should_query_tools(None));
assert!(!should_query_resources(None));
}
/// When `tools` is not advertised, the helper must skip the list call so
/// we don't waste a round trip and pollute the wire log with a request
/// that's destined to return `METHOD_NOT_FOUND`.
#[tokio::test]
async fn query_tools_for_skips_listing_when_capability_not_advertised() {
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let no_caps = caps(false, false);
let result = query_tools_for(Some(&no_caps), "srv", || async move {
calls_clone.fetch_add(1, Ordering::SeqCst);
Ok(vec![test_tool("never")])
})
.await;
assert!(result.is_empty());
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
/// Skips `tools/list` when server info is absent.
#[tokio::test]
async fn query_tools_for_skips_listing_when_server_info_is_none() {
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let result = query_tools_for(None, "srv", || async move {
calls_clone.fetch_add(1, Ordering::SeqCst);
Ok(vec![test_tool("never")])
})
.await;
assert!(result.is_empty());
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
/// Returns listed tools when `tools` is advertised.
#[tokio::test]
async fn query_tools_for_returns_listed_tools_when_capability_advertised() {
let c = caps(true, false);
let expected = vec![test_tool("greet"), test_tool("review")];
let to_return = expected.clone();
let result = query_tools_for(Some(&c), "srv", || async move { Ok(to_return) }).await;
assert_eq!(result, expected);
}
/// Returns an empty vector when the server lists no tools.
#[tokio::test]
async fn query_tools_for_returns_empty_vec_when_server_lists_no_tools() {
let c = caps(true, false);
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let result = query_tools_for(Some(&c), "srv", || async move {
calls_clone.fetch_add(1, Ordering::SeqCst);
Ok(Vec::new())
})
.await;
assert!(result.is_empty());
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
/// **The fail-soft test the bug ticket implicitly demands.** Transport-
/// closed errors must not abort server startup; the helper must log and
/// return an empty vec. This is the regression-protector for #6798's
/// underlying asymmetry — if anyone re-introduces a `return Err(...)` here,
/// this test fails.
#[tokio::test]
async fn query_tools_for_returns_empty_on_transport_error() {
let c = caps(true, false);
let result = query_tools_for(Some(&c), "srv", || async {
Err(rmcp::ServiceError::TransportClosed)
})
.await;
assert!(result.is_empty());
}
/// MCP-protocol errors (e.g. METHOD_NOT_FOUND from a misbehaving server
/// that advertised the capability but rejects the call) also fail soft,
/// so the rest of the server surface still comes up.
#[tokio::test]
async fn query_tools_for_returns_empty_on_mcp_error() {
let c = caps(true, false);
let result = query_tools_for(Some(&c), "srv", || async {
Err(rmcp::ServiceError::McpError(ErrorData {
code: ErrorCode::METHOD_NOT_FOUND,
message: "tools/list not implemented".into(),
data: None,
}))
})
.await;
assert!(result.is_empty());
}
/// Calls the `tools/list` function exactly once per query.
#[tokio::test]
async fn query_tools_for_calls_list_function_exactly_once() {
let c = caps(true, false);
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let _ = query_tools_for(Some(&c), "srv", || async move {
calls_clone.fetch_add(1, Ordering::SeqCst);
Ok(vec![test_tool("p")])
})
.await;
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
/// Keeps the tools-listing decision independent of resource capability state.
#[tokio::test]
async fn query_tools_for_decision_independent_of_other_capabilities() {
let tools = vec![test_tool("x")];
for has_tools in [false, true] {
for has_resources in [false, true] {
let c = caps(has_tools, has_resources);
let to_return = tools.clone();
let result = query_tools_for(Some(&c), "srv", || async move { Ok(to_return) }).await;
if has_tools {
assert_eq!(result, tools);
} else {
assert!(result.is_empty());
}
}
}
}
/// Skips `resources/list` when `resources` is not advertised.
#[tokio::test]
async fn query_resources_for_skips_listing_when_capability_not_advertised() {
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let no_caps = caps(false, false);
let result = query_resources_for(Some(&no_caps), "srv", || async move {
calls_clone.fetch_add(1, Ordering::SeqCst);
Ok(vec![test_resource("file:///nope")])
})
.await;
assert!(result.is_empty());
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
/// Skips `resources/list` when server info is absent.
#[tokio::test]
async fn query_resources_for_skips_listing_when_server_info_is_none() {
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let result = query_resources_for(None, "srv", || async move {
calls_clone.fetch_add(1, Ordering::SeqCst);
Ok(vec![test_resource("file:///nope")])
})
.await;
assert!(result.is_empty());
assert_eq!(calls.load(Ordering::SeqCst), 0);
}
/// Returns listed resources when `resources` is advertised.
#[tokio::test]
async fn query_resources_for_returns_listed_resources_when_capability_advertised() {
let c = caps(false, true);
let expected = vec![test_resource("file:///a"), test_resource("file:///b")];
let to_return = expected.clone();
let result = query_resources_for(Some(&c), "srv", || async move { Ok(to_return) }).await;
assert_eq!(result, expected);
}
/// Fails soft when `resources/list` sees a transport error.
#[tokio::test]
async fn query_resources_for_returns_empty_on_transport_error() {
let c = caps(false, true);
let result = query_resources_for(Some(&c), "srv", || async {
Err(rmcp::ServiceError::TransportClosed)
})
.await;
assert!(result.is_empty());
}
/// Fails soft when `resources/list` returns an MCP protocol error.
#[tokio::test]
async fn query_resources_for_returns_empty_on_mcp_error() {
let c = caps(false, true);
let result = query_resources_for(Some(&c), "srv", || async {
Err(rmcp::ServiceError::McpError(ErrorData {
code: ErrorCode::METHOD_NOT_FOUND,
message: "resources/list not implemented".into(),
data: None,
}))
})
.await;
assert!(result.is_empty());
}
/// Calls the `resources/list` function exactly once per query.
#[tokio::test]
async fn query_resources_for_calls_list_function_exactly_once() {
let c = caps(false, true);
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let _ = query_resources_for(Some(&c), "srv", || async move {
calls_clone.fetch_add(1, Ordering::SeqCst);
Ok(vec![test_resource("file:///a")])
})
.await;
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
crates/mcp/src/sse_transport/auth_impl.rs
+46
View File
@@ -0,0 +1,46 @@
// This file contains code copied from the rmcp crate (https://github.com/modelcontextprotocol/rust-sdk),
// originally located at `crates/rmcp/src/transport/common/auth/sse_client.rs`.
// Used under the terms of the Apache License, Version 2.0.
// See https://github.com/modelcontextprotocol/rust-sdk/blob/main/LICENSE for the full license text.
use http::Uri;
use rmcp::transport::auth::AuthClient;
use super::sse_client::{SseClient, SseTransportError};
impl<C> SseClient for AuthClient<C>
where
C: SseClient,
{
type Error = SseTransportError<C::Error>;
async fn post_message(
&self,
uri: Uri,
message: rmcp::model::ClientJsonRpcMessage,
mut auth_token: Option<String>,
) -> Result<(), SseTransportError<Self::Error>> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.post_message(uri, message, auth_token)
.await
.map_err(SseTransportError::Client)
}
async fn get_stream(
&self,
uri: Uri,
last_event_id: Option<String>,
mut auth_token: Option<String>,
) -> Result<super::client_side_sse::BoxedSseResponse, SseTransportError<Self::Error>> {
if auth_token.is_none() {
auth_token = Some(self.get_access_token().await?);
}
self.http_client
.get_stream(uri, last_event_id, auth_token)
.await
.map_err(SseTransportError::Client)
}
}
crates/mcp/src/sse_transport/client_side_sse.rs
+282
View File
@@ -0,0 +1,282 @@
// This file contains code copied from the rmcp crate (https://github.com/modelcontextprotocol/rust-sdk),
// originally located at `crates/rmcp/src/transport/common/client_side_sse.rs`.
// Used under the terms of the Apache License, Version 2.0.
// See https://github.com/modelcontextprotocol/rust-sdk/blob/main/LICENSE for the full license text.
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{ready, Poll};
use std::time::Duration;
use futures::stream::BoxStream;
use futures::Stream;
use rmcp::model::ServerJsonRpcMessage;
use sse_stream::{Error as SseError, Sse};
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
}
}
/// 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: std::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 },
}
}
}
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 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) => {
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)
}
}
crates/mcp/src/sse_transport/mod.rs
+10
View File
@@ -0,0 +1,10 @@
/// Legacy SSE client transport for MCP, preserved from the rmcp fork after upstream
/// removed SSE transport support in v0.11.0. This allows Warp to continue connecting
/// to MCP servers that only support the older SSE protocol.
mod auth_impl;
mod client_side_sse;
mod reqwest_impl;
mod sse_client;
pub use client_side_sse::{ExponentialBackoff, FixedInterval, NeverRetry, SseRetryPolicy};
pub use sse_client::{SseClient, SseClientConfig, SseClientTransport, SseTransportError};
crates/mcp/src/sse_transport/reqwest_impl.rs
+96
View File
@@ -0,0 +1,96 @@
// This file contains code copied from the rmcp crate (https://github.com/modelcontextprotocol/rust-sdk),
// originally located at `crates/rmcp/src/transport/common/reqwest/sse_client.rs`.
// Used under the terms of the Apache License, Version 2.0.
// See https://github.com/modelcontextprotocol/rust-sdk/blob/main/LICENSE for the full license text.
use std::sync::Arc;
use futures::StreamExt;
use http::Uri;
use reqwest::header::ACCEPT;
use sse_stream::SseStream;
use super::sse_client::{SseClient, SseClientConfig, SseClientTransport, SseTransportError};
const HEADER_LAST_EVENT_ID: &str = "Last-Event-Id";
const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
impl From<reqwest::Error> for SseTransportError<reqwest::Error> {
fn from(e: reqwest::Error) -> Self {
SseTransportError::Client(e)
}
}
impl SseClient for reqwest::Client {
type Error = reqwest::Error;
async fn post_message(
&self,
uri: Uri,
message: rmcp::model::ClientJsonRpcMessage,
auth_token: Option<String>,
) -> Result<(), SseTransportError<Self::Error>> {
let mut request_builder = self.post(uri.to_string()).json(&message);
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
request_builder
.send()
.await
.and_then(|resp| resp.error_for_status())
.map_err(SseTransportError::from)
.map(drop)
}
async fn get_stream(
&self,
uri: Uri,
last_event_id: Option<String>,
auth_token: Option<String>,
) -> Result<super::client_side_sse::BoxedSseResponse, SseTransportError<Self::Error>> {
let mut request_builder = self
.get(uri.to_string())
.header(ACCEPT, EVENT_STREAM_MIME_TYPE);
if let Some(auth_header) = auth_token {
request_builder = request_builder.bearer_auth(auth_header);
}
if let Some(last_event_id) = last_event_id {
request_builder = request_builder.header(HEADER_LAST_EVENT_ID, last_event_id);
}
let response = request_builder.send().await?;
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()) {
return Err(SseTransportError::UnexpectedContentType(Some(
String::from_utf8_lossy(ct.as_bytes()).to_string(),
)));
}
}
None => {
return Err(SseTransportError::UnexpectedContentType(None));
}
}
let event_stream = SseStream::from_byte_stream(response.bytes_stream()).boxed();
Ok(event_stream)
}
}
impl SseClientTransport<reqwest::Client> {
/// Creates a new transport using reqwest with the specified SSE endpoint.
///
/// This is a convenience method that creates a transport using the default
/// reqwest client.
pub async fn start(
uri: impl Into<Arc<str>>,
) -> Result<Self, SseTransportError<reqwest::Error>> {
SseClientTransport::start_with_client(
reqwest::Client::default(),
SseClientConfig {
sse_endpoint: uri.into(),
..Default::default()
},
)
.await
}
}
crates/mcp/src/sse_transport/sse_client.rs
+260
View File
@@ -0,0 +1,260 @@
// This file contains code copied from the rmcp crate (https://github.com/modelcontextprotocol/rust-sdk),
// originally located at `crates/rmcp/src/transport/sse_client.rs`.
// Used under the terms of the Apache License, Version 2.0.
// See https://github.com/modelcontextprotocol/rust-sdk/blob/main/LICENSE for the full license text.
//
// Reference: <https://html.spec.whatwg.org/multipage/server-sent-events.html>
use std::pin::Pin;
use std::sync::{Arc, RwLock};
use futures::future::BoxFuture;
use futures::StreamExt;
use http::Uri;
use rmcp::model::{ClientJsonRpcMessage, ServerJsonRpcMessage};
use rmcp::transport::Transport;
use rmcp::RoleClient;
use sse_stream::{Error as SseError, Sse};
use thiserror::Error;
use super::client_side_sse::{
BoxedSseResponse, SseAutoReconnectStream, SseRetryPolicy, SseStreamReconnect,
};
#[derive(Error, Debug)]
pub enum SseTransportError<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 content type: {0:?}")]
UnexpectedContentType(Option<String>),
#[error("Auth error: {0}")]
Auth(#[from] rmcp::transport::AuthError),
#[error("Invalid uri: {0}")]
InvalidUri(#[from] http::uri::InvalidUri),
#[error("Invalid uri parts: {0}")]
InvalidUriParts(#[from] http::uri::InvalidUriParts),
}
pub trait SseClient: Clone + Send + Sync + 'static {
type Error: std::error::Error + Send + Sync + 'static;
fn post_message(
&self,
uri: Uri,
message: ClientJsonRpcMessage,
auth_token: Option<String>,
) -> impl std::future::Future<Output = Result<(), SseTransportError<Self::Error>>> + Send + '_;
fn get_stream(
&self,
uri: Uri,
last_event_id: Option<String>,
auth_token: Option<String>,
) -> impl std::future::Future<Output = Result<BoxedSseResponse, SseTransportError<Self::Error>>>
+ Send
+ '_;
}
/// Helper that refreshes the POST endpoint whenever the server emits
/// control frames during SSE reconnect; used together with
/// [`SseAutoReconnectStream`].
struct SseClientReconnect<C> {
pub client: C,
pub uri: Uri,
pub message_endpoint: Arc<RwLock<Uri>>,
}
impl<C: SseClient> SseStreamReconnect for SseClientReconnect<C> {
type Error = SseTransportError<C::Error>;
type Future = BoxFuture<'static, Result<BoxedSseResponse, 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 last_event_id = last_event_id.map(|s| s.to_owned());
Box::pin(async move { client.get_stream(uri, last_event_id, None).await })
}
fn handle_control_event(&mut self, event: &Sse) -> Result<(), Self::Error> {
if event.event.as_deref() != Some("endpoint") {
return Ok(());
}
let Some(data) = event.data.as_ref() else {
return Ok(());
};
// Servers typically resend the message POST endpoint (often with a new
// sessionId) when a stream reconnects. Reuse `message_endpoint` helper
// to resolve it and update the shared URI.
let new_endpoint = message_endpoint(self.uri.clone(), data.clone())
.map_err(SseTransportError::InvalidUri)?;
*self
.message_endpoint
.write()
.expect("message endpoint lock poisoned") = new_endpoint;
Ok(())
}
fn handle_stream_error(
&mut self,
error: &(dyn std::error::Error + 'static),
last_event_id: Option<&str>,
) {
tracing::warn!(
uri = %self.uri,
last_event_id = last_event_id.unwrap_or(""),
"sse stream error: {error}"
);
}
}
type ServerMessageStream<C> = Pin<Box<SseAutoReconnectStream<SseClientReconnect<C>>>>;
/// A client-agnostic SSE transport for MCP that supports Server-Sent Events.
///
/// This transport allows you to choose your preferred HTTP client implementation
/// by implementing the [`SseClient`] trait. The transport handles SSE streaming
/// and automatic reconnection.
pub struct SseClientTransport<C: SseClient> {
client: C,
config: SseClientConfig,
/// Current POST endpoint; refreshed when the server sends new endpoint
/// control frames.
message_endpoint: Arc<RwLock<Uri>>,
stream: Option<ServerMessageStream<C>>,
}
impl<C: SseClient> Transport<RoleClient> for SseClientTransport<C> {
type Error = SseTransportError<C::Error>;
async fn receive(&mut self) -> Option<ServerJsonRpcMessage> {
self.stream.as_mut()?.next().await?.ok()
}
fn send(
&mut self,
item: rmcp::service::TxJsonRpcMessage<RoleClient>,
) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send + 'static {
let client = self.client.clone();
let message_endpoint = self.message_endpoint.clone();
async move {
let uri = {
let guard = message_endpoint
.read()
.expect("message endpoint lock poisoned");
guard.clone()
};
client.post_message(uri, item, None).await
}
}
async fn close(&mut self) -> Result<(), Self::Error> {
self.stream.take();
Ok(())
}
}
impl<C: SseClient + std::fmt::Debug> std::fmt::Debug for SseClientTransport<C> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SseClientTransport")
.field("client", &self.client)
.field("config", &self.config)
.finish()
}
}
impl<C: SseClient> SseClientTransport<C> {
pub async fn start_with_client(
client: C,
config: SseClientConfig,
) -> Result<Self, SseTransportError<C::Error>> {
let sse_endpoint = config.sse_endpoint.as_ref().parse::<http::Uri>()?;
let mut sse_stream = client.get_stream(sse_endpoint.clone(), None, None).await?;
let initial_message_endpoint = if let Some(endpoint) = config.use_message_endpoint.clone() {
let ep = endpoint.parse::<http::Uri>()?;
let mut sse_endpoint_parts = sse_endpoint.clone().into_parts();
sse_endpoint_parts.path_and_query = ep.into_parts().path_and_query;
Uri::from_parts(sse_endpoint_parts)?
} else {
// Wait for the endpoint event.
loop {
let sse = sse_stream
.next()
.await
.ok_or(SseTransportError::UnexpectedEndOfStream)??;
let Some("endpoint") = sse.event.as_deref() else {
continue;
};
let ep = sse.data.unwrap_or_default();
break message_endpoint(sse_endpoint.clone(), ep)?;
}
};
let message_endpoint = Arc::new(RwLock::new(initial_message_endpoint));
let stream = Box::pin(SseAutoReconnectStream::new(
sse_stream,
SseClientReconnect {
client: client.clone(),
uri: sse_endpoint.clone(),
message_endpoint: message_endpoint.clone(),
},
config.retry_policy.clone(),
));
Ok(Self {
client,
config,
message_endpoint,
stream: Some(stream),
})
}
}
fn message_endpoint(base: http::Uri, endpoint: String) -> Result<http::Uri, http::uri::InvalidUri> {
// If endpoint is a full URL, parse and return it directly.
if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
return endpoint.parse::<http::Uri>();
}
let mut base_parts = base.into_parts();
let endpoint_clone = endpoint.clone();
if endpoint.starts_with("?") {
// Query only - keep base path and append query.
if let Some(base_path_and_query) = &base_parts.path_and_query {
let base_path = base_path_and_query.path();
base_parts.path_and_query = Some(format!("{base_path}{endpoint}").parse()?);
} else {
base_parts.path_and_query = Some(format!("/{endpoint}").parse()?);
}
} else {
// Path (with optional query) - replace entire path_and_query.
let path_to_use = if endpoint.starts_with("/") {
endpoint // Use absolute path as-is.
} else {
format!("/{endpoint}") // Make relative path absolute.
};
base_parts.path_and_query = Some(path_to_use.parse()?);
}
http::Uri::from_parts(base_parts).map_err(|_| endpoint_clone.parse::<http::Uri>().unwrap_err())
}
#[derive(Debug, Clone)]
pub struct SseClientConfig {
/// The client SSE endpoint URL.
pub sse_endpoint: Arc<str>,
pub retry_policy: Arc<dyn SseRetryPolicy>,
/// If this is set, the client will use this endpoint to send messages and
/// skip waiting for the endpoint event.
pub use_message_endpoint: Option<String>,
}
impl Default for SseClientConfig {
fn default() -> Self {
Self {
sse_endpoint: "".into(),
retry_policy: Arc::new(super::client_side_sse::FixedInterval::default()),
use_message_endpoint: None,
}
}
}