feat: add modern client lifecycle modes (SEP-2575) (#995)
* feat: add modern client lifecycle modes * refactor: drop modern naming and dead cleanup * fix: gate request metadata on inline lifecycle only
This commit is contained in:
Dale Seo
GitHub
parent
81745f1375
commit
2e2c79173c
7 changed files with 1525 additions and 306 deletions
|
|
@ -41,22 +41,30 @@ impl<H: ServerHandler> Service<RoleServer> for H {
|
|||
));
|
||||
}
|
||||
}
|
||||
if matches!(&request, ClientRequest::DiscoverRequest(_)) {
|
||||
if requested_version.is_none() {
|
||||
// Self-contained metadata is required only when the request itself uses
|
||||
// the inline lifecycle: a discover opener, a session that started without
|
||||
// `initialize`, or a request that declares 2026-07-28+ in its own _meta.
|
||||
// Sessions that negotiated via `initialize` (or `serve_directly`) keep the
|
||||
// session model and may omit per-request metadata.
|
||||
let requires_request_metadata = uses_inline_negotiation
|
||||
&& (matches!(&request, ClientRequest::DiscoverRequest(_))
|
||||
|| context.peer.request_metadata_required()
|
||||
|| requested_version.as_ref().is_some_and(|version| {
|
||||
version.as_str() >= ProtocolVersion::V_2026_07_28.as_str()
|
||||
}));
|
||||
if requires_request_metadata {
|
||||
// Inline lifecycle requests are defined by the 2026-07-28 protocol.
|
||||
// Validate that lifecycle contract even when a request selects an
|
||||
// older application protocol version.
|
||||
let missing = context
|
||||
.meta
|
||||
.missing_required_keys(&ProtocolVersion::V_2026_07_28);
|
||||
if !missing.is_empty() {
|
||||
return Err(McpError::invalid_params(
|
||||
"server/discover requires protocolVersion in request _meta",
|
||||
None,
|
||||
));
|
||||
}
|
||||
if context.meta.client_info().is_none() {
|
||||
return Err(McpError::invalid_params(
|
||||
"server/discover requires clientInfo in request _meta",
|
||||
None,
|
||||
));
|
||||
}
|
||||
if context.meta.client_capabilities().is_none() {
|
||||
return Err(McpError::invalid_params(
|
||||
"server/discover requires clientCapabilities in request _meta",
|
||||
format!(
|
||||
"request _meta is missing or has malformed required fields: {}",
|
||||
missing.join(", ")
|
||||
),
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,10 +16,13 @@ pub use handler::client::ClientHandler;
|
|||
pub use handler::server::ServerHandler;
|
||||
#[cfg(feature = "server")]
|
||||
pub use handler::server::wrapper::Json;
|
||||
#[cfg(feature = "client")]
|
||||
pub use service::{
|
||||
ClientLifecycleMode, ClientServiceExt, RoleClient, select_protocol_version, serve_client,
|
||||
serve_client_with_lifecycle,
|
||||
};
|
||||
#[cfg(any(feature = "client", feature = "server"))]
|
||||
pub use service::{Peer, Service, ServiceError, ServiceExt};
|
||||
#[cfg(feature = "client")]
|
||||
pub use service::{RoleClient, select_protocol_version, serve_client};
|
||||
#[cfg(feature = "server")]
|
||||
pub use service::{RoleServer, serve_server};
|
||||
|
||||
|
|
|
|||
|
|
@ -439,6 +439,19 @@ impl RequestMetaObject {
|
|||
meta
|
||||
}
|
||||
|
||||
/// Create request metadata with the client context SEP-2575 requires on every request.
|
||||
pub fn with_client_context(
|
||||
protocol_version: ProtocolVersion,
|
||||
client_info: Implementation,
|
||||
client_capabilities: ClientCapabilities,
|
||||
) -> Self {
|
||||
let mut meta = Self::new();
|
||||
meta.set_protocol_version(protocol_version);
|
||||
meta.set_client_info(client_info);
|
||||
meta.set_client_capabilities(client_capabilities);
|
||||
meta
|
||||
}
|
||||
|
||||
pub(crate) fn static_empty() -> &'static Self {
|
||||
static EMPTY: std::sync::OnceLock<RequestMetaObject> = std::sync::OnceLock::new();
|
||||
EMPTY.get_or_init(Default::default)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
use std::sync::OnceLock;
|
||||
|
||||
use futures::FutureExt;
|
||||
#[cfg(not(feature = "local"))]
|
||||
use futures::future::BoxFuture;
|
||||
|
|
@ -51,9 +53,10 @@ use crate::model::ServerNotification;
|
|||
use crate::{
|
||||
error::ErrorData as McpError,
|
||||
model::{
|
||||
CancelledNotification, CancelledNotificationParam, Extensions, GetExtensions, GetMeta,
|
||||
JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse,
|
||||
NotificationMetaObject, NumberOrString, ProgressToken, RequestId, RequestMetaObject,
|
||||
CancelledNotification, CancelledNotificationParam, ClientCapabilities, Extensions,
|
||||
GetExtensions, GetMeta, Implementation, JsonRpcError, JsonRpcMessage, JsonRpcNotification,
|
||||
JsonRpcRequest, JsonRpcResponse, NotificationMetaObject, NumberOrString, ProgressToken,
|
||||
ProtocolVersion, RequestId, RequestMetaObject,
|
||||
},
|
||||
transport::{DynamicTransportError, IntoTransport, Transport},
|
||||
};
|
||||
|
|
@ -515,6 +518,13 @@ pub(crate) enum PeerSinkMessage<R: ServiceRole> {
|
|||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct ClientRequestMetadata {
|
||||
pub protocol_version: ProtocolVersion,
|
||||
pub client_info: Implementation,
|
||||
pub client_capabilities: ClientCapabilities,
|
||||
}
|
||||
|
||||
/// 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.
|
||||
|
|
@ -527,6 +537,8 @@ pub struct Peer<R: ServiceRole> {
|
|||
progress_token_provider: Arc<dyn ProgressTokenProvider>,
|
||||
progress_timeout_watchers: ProgressTimeoutWatchers,
|
||||
info: Arc<std::sync::RwLock<Option<Arc<R::PeerInfo>>>>,
|
||||
client_request_metadata: Arc<OnceLock<ClientRequestMetadata>>,
|
||||
request_metadata_required: Arc<std::sync::atomic::AtomicBool>,
|
||||
}
|
||||
|
||||
impl<R: ServiceRole> std::fmt::Debug for Peer<R> {
|
||||
|
|
@ -563,6 +575,14 @@ impl PeerRequestOptions {
|
|||
}
|
||||
}
|
||||
|
||||
/// Adds request metadata while preserving any other configured options.
|
||||
///
|
||||
/// Explicit values take precedence over discover-lifecycle metadata defaults.
|
||||
pub fn with_meta(mut self, meta: RequestMetaObject) -> Self {
|
||||
self.meta = Some(meta);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn reset_timeout_on_progress(mut self) -> Self {
|
||||
self.reset_timeout_on_progress = true;
|
||||
self
|
||||
|
|
@ -588,6 +608,8 @@ impl<R: ServiceRole> Peer<R> {
|
|||
progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()),
|
||||
progress_timeout_watchers: Default::default(),
|
||||
info: Arc::new(std::sync::RwLock::new(peer_info.map(Arc::new))),
|
||||
client_request_metadata: Default::default(),
|
||||
request_metadata_required: Default::default(),
|
||||
},
|
||||
rx,
|
||||
)
|
||||
|
|
@ -625,6 +647,12 @@ impl<R: ServiceRole> Peer<R> {
|
|||
) -> Result<RequestHandle<R>, ServiceError> {
|
||||
let id = self.request_id_provider.next_request_id();
|
||||
let progress_token = self.progress_token_provider.next_progress_token();
|
||||
if let Some(metadata) = self.client_request_metadata.get() {
|
||||
let meta = request.get_meta_mut();
|
||||
meta.set_protocol_version(metadata.protocol_version.clone());
|
||||
meta.set_client_info(metadata.client_info.clone());
|
||||
meta.set_client_capabilities(metadata.client_capabilities.clone());
|
||||
}
|
||||
if let Some(meta) = options.meta.clone() {
|
||||
request.get_meta_mut().extend(meta);
|
||||
}
|
||||
|
|
@ -703,6 +731,21 @@ impl<R: ServiceRole> Peer<R> {
|
|||
*self.info.write().expect("peer info lock poisoned") = Some(Arc::new(info));
|
||||
}
|
||||
|
||||
pub(crate) fn set_client_request_metadata(&self, metadata: ClientRequestMetadata) {
|
||||
let result = self.client_request_metadata.set(metadata);
|
||||
debug_assert!(result.is_ok(), "client request metadata set more than once");
|
||||
}
|
||||
|
||||
pub(crate) fn require_request_metadata(&self) {
|
||||
self.request_metadata_required
|
||||
.store(true, std::sync::atomic::Ordering::Release);
|
||||
}
|
||||
|
||||
pub(crate) fn request_metadata_required(&self) -> bool {
|
||||
self.request_metadata_required
|
||||
.load(std::sync::atomic::Ordering::Acquire)
|
||||
}
|
||||
|
||||
pub fn is_transport_closed(&self) -> bool {
|
||||
self.tx.is_closed()
|
||||
}
|
||||
|
|
@ -883,11 +926,35 @@ impl<R: ServiceRole> RequestContext<R> {
|
|||
|
||||
#[cfg(feature = "server")]
|
||||
impl RequestContext<RoleServer> {
|
||||
/// The protocol version the client negotiated, or `None` before peer info is recorded.
|
||||
/// The current request's protocol version, falling back to legacy handshake state.
|
||||
pub fn protocol_version(&self) -> Option<crate::model::ProtocolVersion> {
|
||||
self.peer
|
||||
.peer_info()
|
||||
.map(|info| info.protocol_version.clone())
|
||||
self.meta.protocol_version().or_else(|| {
|
||||
self.peer
|
||||
.peer_info()
|
||||
.map(|info| info.protocol_version.clone())
|
||||
})
|
||||
}
|
||||
|
||||
/// The current request's client implementation, falling back only for legacy sessions.
|
||||
pub fn client_info(&self) -> Option<Implementation> {
|
||||
if self.peer.request_metadata_required() {
|
||||
self.meta.client_info()
|
||||
} else {
|
||||
self.meta
|
||||
.client_info()
|
||||
.or_else(|| self.peer.peer_info().map(|info| info.client_info.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
/// The current request's client capabilities, falling back only for legacy sessions.
|
||||
pub fn client_capabilities(&self) -> Option<ClientCapabilities> {
|
||||
if self.peer.request_metadata_required() {
|
||||
self.meta.client_capabilities()
|
||||
} else {
|
||||
self.meta
|
||||
.client_capabilities()
|
||||
.or_else(|| self.peer.peer_info().map(|info| info.capabilities.clone()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -54,6 +54,17 @@ pub enum ClientInitializeError {
|
|||
#[error("JSON-RPC error: {0}")]
|
||||
JsonRpcError(ErrorData),
|
||||
|
||||
#[error(
|
||||
"no compatible protocol version (client: {client_supported:?}, server: {server_supported:?})"
|
||||
)]
|
||||
NoCompatibleProtocolVersion {
|
||||
client_supported: Vec<ProtocolVersion>,
|
||||
server_supported: Vec<ProtocolVersion>,
|
||||
},
|
||||
|
||||
#[error("discover startup requires at least one preferred protocol version")]
|
||||
NoPreferredProtocolVersion,
|
||||
|
||||
#[error("Cancelled")]
|
||||
Cancelled,
|
||||
}
|
||||
|
|
@ -177,6 +188,43 @@ impl ServiceRole for RoleClient {
|
|||
|
||||
pub type ServerSink = Peer<RoleClient>;
|
||||
|
||||
/// Selects how a client establishes its MCP lifecycle.
|
||||
///
|
||||
/// Existing [`ServiceExt::serve`] behavior remains legacy initialization.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
#[non_exhaustive]
|
||||
pub enum ClientLifecycleMode {
|
||||
/// Use the legacy `initialize` / `notifications/initialized` handshake.
|
||||
Initialize,
|
||||
/// Use `server/discover` and send self-contained per-request metadata.
|
||||
Discover {
|
||||
preferred_versions: Vec<ProtocolVersion>,
|
||||
},
|
||||
/// Probe with `server/discover`, falling back only when the peer proves it is legacy.
|
||||
Auto {
|
||||
preferred_versions: Vec<ProtocolVersion>,
|
||||
legacy_version: Option<ProtocolVersion>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Client-specific lifecycle entry points.
|
||||
pub trait ClientServiceExt: Service<RoleClient> + Sized {
|
||||
fn serve_with_lifecycle<T, E, A>(
|
||||
self,
|
||||
transport: T,
|
||||
lifecycle: ClientLifecycleMode,
|
||||
) -> impl Future<Output = Result<RunningService<RoleClient, Self>, ClientInitializeError>>
|
||||
+ MaybeSendFuture
|
||||
where
|
||||
T: IntoTransport<RoleClient, E, A>,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
serve_client_with_lifecycle(self, transport, lifecycle)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: Service<RoleClient>> ClientServiceExt for S {}
|
||||
|
||||
impl<S: Service<RoleClient>> ServiceExt<RoleClient> for S {
|
||||
fn serve_with_ct<T, E, A>(
|
||||
self,
|
||||
|
|
@ -202,7 +250,13 @@ where
|
|||
T: IntoTransport<RoleClient, E, A>,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
serve_client_with_ct(service, transport, Default::default()).await
|
||||
serve_client_with_lifecycle_and_ct(
|
||||
service,
|
||||
transport,
|
||||
ClientLifecycleMode::Initialize,
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn serve_client_with_ct<S, T, E, A>(
|
||||
|
|
@ -210,13 +264,41 @@ pub async fn serve_client_with_ct<S, T, E, A>(
|
|||
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,
|
||||
{
|
||||
serve_client_with_lifecycle_and_ct(service, transport, ClientLifecycleMode::Initialize, ct)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn serve_client_with_lifecycle<S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
lifecycle: ClientLifecycleMode,
|
||||
) -> Result<RunningService<RoleClient, S>, ClientInitializeError>
|
||||
where
|
||||
S: Service<RoleClient>,
|
||||
T: IntoTransport<RoleClient, E, A>,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
serve_client_with_lifecycle_and_ct(service, transport, lifecycle, Default::default()).await
|
||||
}
|
||||
|
||||
pub async fn serve_client_with_lifecycle_and_ct<S, T, E, A>(
|
||||
service: S,
|
||||
transport: T,
|
||||
lifecycle: ClientLifecycleMode,
|
||||
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 }
|
||||
result = serve_client_with_ct_inner(service, transport.into_transport(), lifecycle, ct.clone()) => { result }
|
||||
_ = ct.cancelled() => {
|
||||
Err(ClientInitializeError::Cancelled)
|
||||
}
|
||||
|
|
@ -226,6 +308,7 @@ where
|
|||
async fn serve_client_with_ct_inner<S, T>(
|
||||
service: S,
|
||||
transport: T,
|
||||
lifecycle: ClientLifecycleMode,
|
||||
ct: CancellationToken,
|
||||
) -> Result<RunningService<RoleClient, S>, ClientInitializeError>
|
||||
where
|
||||
|
|
@ -234,12 +317,71 @@ where
|
|||
{
|
||||
let mut transport = transport.into_transport();
|
||||
let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
|
||||
let (peer, peer_rx) = Peer::new(id_provider.clone(), None);
|
||||
let client_info = service.get_info();
|
||||
|
||||
// service
|
||||
match lifecycle {
|
||||
ClientLifecycleMode::Initialize => {
|
||||
legacy_startup(&service, &mut transport, &id_provider, &peer, client_info).await?;
|
||||
}
|
||||
ClientLifecycleMode::Discover { preferred_versions } => {
|
||||
discover_startup(
|
||||
&service,
|
||||
&mut transport,
|
||||
&id_provider,
|
||||
&peer,
|
||||
&client_info,
|
||||
preferred_versions,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
ClientLifecycleMode::Auto {
|
||||
preferred_versions,
|
||||
legacy_version,
|
||||
} => {
|
||||
let discover_result = discover_startup(
|
||||
&service,
|
||||
&mut transport,
|
||||
&id_provider,
|
||||
&peer,
|
||||
&client_info,
|
||||
preferred_versions,
|
||||
)
|
||||
.await;
|
||||
match discover_result {
|
||||
Ok(()) => {}
|
||||
Err(ClientInitializeError::JsonRpcError(error))
|
||||
if error.code == crate::model::ErrorCode::METHOD_NOT_FOUND =>
|
||||
{
|
||||
let mut legacy_info = client_info;
|
||||
if let Some(version) = legacy_version {
|
||||
legacy_info.protocol_version = version;
|
||||
}
|
||||
legacy_startup(&service, &mut transport, &id_provider, &peer, legacy_info)
|
||||
.await?;
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(serve_inner(service, transport, peer, peer_rx, ct))
|
||||
}
|
||||
|
||||
async fn legacy_startup<S, T>(
|
||||
service: &S,
|
||||
transport: &mut T,
|
||||
id_provider: &Arc<AtomicU32RequestIdProvider>,
|
||||
peer: &Peer<RoleClient>,
|
||||
client_info: ClientInfo,
|
||||
) -> Result<(), ClientInitializeError>
|
||||
where
|
||||
S: Service<RoleClient>,
|
||||
T: Transport<RoleClient> + 'static,
|
||||
{
|
||||
let id = id_provider.next_request_id();
|
||||
let init_request = InitializeRequest {
|
||||
method: Default::default(),
|
||||
params: service.get_info(),
|
||||
params: client_info,
|
||||
extensions: Default::default(),
|
||||
};
|
||||
transport
|
||||
|
|
@ -253,15 +395,8 @@ where
|
|||
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?;
|
||||
let (response, response_id) =
|
||||
expect_response(transport, "initialize response", service, peer.clone()).await?;
|
||||
|
||||
if id != response_id {
|
||||
return Err(ClientInitializeError::ConflictInitResponseId(
|
||||
|
|
@ -285,7 +420,115 @@ where
|
|||
transport.send(notification).await.map_err(|error| {
|
||||
ClientInitializeError::transport::<T>(error, "send initialized notification")
|
||||
})?;
|
||||
Ok(serve_inner(service, transport, peer, peer_rx, ct))
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn discover_startup<S, T>(
|
||||
service: &S,
|
||||
transport: &mut T,
|
||||
id_provider: &Arc<AtomicU32RequestIdProvider>,
|
||||
peer: &Peer<RoleClient>,
|
||||
client_info: &ClientInfo,
|
||||
preferred_versions: Vec<ProtocolVersion>,
|
||||
) -> Result<(), ClientInitializeError>
|
||||
where
|
||||
S: Service<RoleClient>,
|
||||
T: Transport<RoleClient> + 'static,
|
||||
{
|
||||
if preferred_versions.is_empty() {
|
||||
return Err(ClientInitializeError::NoPreferredProtocolVersion);
|
||||
}
|
||||
|
||||
let mut attempted = Vec::new();
|
||||
let mut candidate = preferred_versions[0].clone();
|
||||
loop {
|
||||
attempted.push(candidate.clone());
|
||||
|
||||
let meta = RequestMetaObject::with_client_context(
|
||||
candidate.clone(),
|
||||
client_info.client_info.clone(),
|
||||
client_info.capabilities.clone(),
|
||||
);
|
||||
let mut discover = DiscoverRequest::new(DiscoverRequestParams {});
|
||||
discover.extensions.insert(meta);
|
||||
let id = id_provider.next_request_id();
|
||||
transport
|
||||
.send(ClientJsonRpcMessage::request(
|
||||
ClientRequest::DiscoverRequest(discover),
|
||||
id.clone(),
|
||||
))
|
||||
.await
|
||||
.map_err(|error| {
|
||||
ClientInitializeError::transport::<T>(error, "send discover request")
|
||||
})?;
|
||||
|
||||
match expect_response(transport, "discover response", service, peer.clone()).await {
|
||||
Ok((ServerResult::DiscoverResult(result), response_id)) => {
|
||||
if response_id != id {
|
||||
return Err(ClientInitializeError::ConflictInitResponseId(
|
||||
id,
|
||||
response_id,
|
||||
));
|
||||
}
|
||||
let Some(selected) =
|
||||
select_protocol_version(&preferred_versions, &result.supported_versions)
|
||||
else {
|
||||
return Err(ClientInitializeError::NoCompatibleProtocolVersion {
|
||||
client_supported: preferred_versions,
|
||||
server_supported: result.supported_versions,
|
||||
});
|
||||
};
|
||||
peer.set_peer_info(ServerInfo {
|
||||
protocol_version: selected.clone(),
|
||||
capabilities: result.capabilities,
|
||||
server_info: result.server_info,
|
||||
instructions: result.instructions,
|
||||
meta: result.meta,
|
||||
});
|
||||
peer.set_client_request_metadata(ClientRequestMetadata {
|
||||
protocol_version: selected,
|
||||
client_info: client_info.client_info.clone(),
|
||||
client_capabilities: client_info.capabilities.clone(),
|
||||
});
|
||||
return Ok(());
|
||||
}
|
||||
Ok((response, _)) => {
|
||||
return Err(ClientInitializeError::ExpectedInitResult(Some(response)));
|
||||
}
|
||||
Err(ClientInitializeError::JsonRpcError(error))
|
||||
if error.code == crate::model::ErrorCode::UNSUPPORTED_PROTOCOL_VERSION =>
|
||||
{
|
||||
let supported = error
|
||||
.data
|
||||
.as_ref()
|
||||
.and_then(|data| data.get("supported"))
|
||||
.cloned()
|
||||
.and_then(|value| serde_json::from_value::<Vec<ProtocolVersion>>(value).ok())
|
||||
.unwrap_or_default();
|
||||
let may_retry_current = attempted
|
||||
.iter()
|
||||
.filter(|version| *version == &candidate)
|
||||
.count()
|
||||
== 1;
|
||||
let next = preferred_versions
|
||||
.iter()
|
||||
.find(|version| {
|
||||
supported.contains(version)
|
||||
&& (!attempted.contains(version)
|
||||
|| (may_retry_current && *version == &candidate))
|
||||
})
|
||||
.cloned();
|
||||
let Some(next) = next else {
|
||||
return Err(ClientInitializeError::NoCompatibleProtocolVersion {
|
||||
client_supported: preferred_versions,
|
||||
server_supported: supported,
|
||||
});
|
||||
};
|
||||
candidate = next;
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! method {
|
||||
|
|
|
|||
|
|
@ -217,12 +217,41 @@ where
|
|||
}
|
||||
};
|
||||
|
||||
let ClientRequest::InitializeRequest(peer_info) = &request else {
|
||||
return Err(ServerInitializeError::ExpectedInitializeRequest(Some(
|
||||
ClientJsonRpcMessage::request(request, id),
|
||||
)));
|
||||
let initialize_request = match request {
|
||||
ClientRequest::InitializeRequest(request) => request,
|
||||
mut request => {
|
||||
if !request
|
||||
.get_meta()
|
||||
.missing_required_keys(&ProtocolVersion::V_2026_07_28)
|
||||
.is_empty()
|
||||
{
|
||||
return Err(ServerInitializeError::ExpectedInitializeRequest(Some(
|
||||
ClientJsonRpcMessage::request(request, id),
|
||||
)));
|
||||
}
|
||||
let (peer, peer_rx) = Peer::new(id_provider, None);
|
||||
peer.require_request_metadata();
|
||||
let context = RequestContext {
|
||||
ct: ct.child_token(),
|
||||
id: id.clone(),
|
||||
meta: std::mem::take(request.get_meta_mut()),
|
||||
extensions: std::mem::take(request.extensions_mut()),
|
||||
peer: peer.clone(),
|
||||
};
|
||||
let response = match service.handle_request(request, context).await {
|
||||
Ok(result) => ServerJsonRpcMessage::response(result, id),
|
||||
Err(error) => ServerJsonRpcMessage::error(error, Some(id)),
|
||||
};
|
||||
transport.send(response).await.map_err(|error| {
|
||||
ServerInitializeError::transport::<T>(error, "sending negotiated request response")
|
||||
})?;
|
||||
return Ok(serve_inner(service, transport, peer, peer_rx, ct));
|
||||
}
|
||||
};
|
||||
let (peer, peer_rx) = Peer::new(id_provider, Some(peer_info.params.clone()));
|
||||
let requested_protocol_version = initialize_request.params.protocol_version.clone();
|
||||
let mut negotiated_peer_info = initialize_request.params.clone();
|
||||
let (peer, peer_rx) = Peer::new(id_provider, Some(negotiated_peer_info.clone()));
|
||||
let request = ClientRequest::InitializeRequest(initialize_request);
|
||||
let context = RequestContext {
|
||||
ct: ct.child_token(),
|
||||
id: id.clone(),
|
||||
|
|
@ -231,7 +260,7 @@ where
|
|||
peer: peer.clone(),
|
||||
};
|
||||
// Send initialize response
|
||||
let init_response = service.handle_request(request.clone(), context).await;
|
||||
let init_response = service.handle_request(request, context).await;
|
||||
let mut init_response = match init_response {
|
||||
Ok(ServerResult::InitializeResult(init_response)) => init_response,
|
||||
Ok(result) => {
|
||||
|
|
@ -247,13 +276,10 @@ where
|
|||
return Err(ServerInitializeError::InitializeFailed(e));
|
||||
}
|
||||
};
|
||||
init_response.protocol_version = negotiate_protocol_version(
|
||||
&peer_info.params.protocol_version,
|
||||
init_response.protocol_version,
|
||||
);
|
||||
init_response.protocol_version =
|
||||
negotiate_protocol_version(&requested_protocol_version, init_response.protocol_version);
|
||||
// Update peer_info so context.protocol_version() reflects the negotiated
|
||||
// version in all subsequent request handlers.
|
||||
let mut negotiated_peer_info = peer_info.params.clone();
|
||||
negotiated_peer_info.protocol_version = init_response.protocol_version.clone();
|
||||
peer.set_peer_info(negotiated_peer_info);
|
||||
transport
|
||||
|
|
|
|||
|
|
@ -19,8 +19,9 @@ use super::common::client_side_sse::{
|
|||
use crate::{
|
||||
RoleClient,
|
||||
model::{
|
||||
ClientJsonRpcMessage, ClientNotification, ErrorData, InitializedNotification, JsonObject,
|
||||
ProtocolVersion, RequestId, ServerJsonRpcMessage, ServerResult,
|
||||
ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetMeta,
|
||||
InitializedNotification, JsonObject, ProtocolVersion, RequestId, ServerJsonRpcMessage,
|
||||
ServerResult,
|
||||
},
|
||||
transport::{
|
||||
common::{client_side_sse::SseAutoReconnectStream, mcp_headers},
|
||||
|
|
@ -29,6 +30,7 @@ use crate::{
|
|||
};
|
||||
|
||||
type BoxedSseStream = BoxStream<'static, Result<Sse, SseError>>;
|
||||
const SESSION_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
fn build_request_headers(
|
||||
base: &HashMap<HeaderName, HeaderValue>,
|
||||
|
|
@ -58,6 +60,27 @@ fn build_request_headers(
|
|||
headers
|
||||
}
|
||||
|
||||
fn request_version_headers(
|
||||
base: &HashMap<HeaderName, HeaderValue>,
|
||||
message: &ClientJsonRpcMessage,
|
||||
fallback: &ProtocolVersion,
|
||||
tool_cache: &HashMap<String, Arc<JsonObject>>,
|
||||
) -> (ProtocolVersion, HashMap<HeaderName, HeaderValue>) {
|
||||
let version = match message {
|
||||
ClientJsonRpcMessage::Request(request) => request
|
||||
.request
|
||||
.get_meta()
|
||||
.protocol_version()
|
||||
.unwrap_or_else(|| fallback.clone()),
|
||||
_ => fallback.clone(),
|
||||
};
|
||||
let mut headers = build_request_headers(base, message, tool_cache, &version);
|
||||
if let Ok(value) = HeaderValue::from_str(version.as_str()) {
|
||||
headers.insert(HeaderName::from_static("mcp-protocol-version"), value);
|
||||
}
|
||||
(version, headers)
|
||||
}
|
||||
|
||||
fn cache_tools_from_response(
|
||||
cache: &mut HashMap<String, Arc<JsonObject>>,
|
||||
message: &ServerJsonRpcMessage,
|
||||
|
|
@ -184,7 +207,10 @@ pub enum StreamableHttpProtocolError {
|
|||
MissingSessionIdInResponse,
|
||||
}
|
||||
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
#[expect(
|
||||
clippy::large_enum_variant,
|
||||
reason = "boxing the streaming response would add an allocation to the common response path"
|
||||
)]
|
||||
#[non_exhaustive]
|
||||
pub enum StreamableHttpPostResponse {
|
||||
Accepted,
|
||||
|
|
@ -651,6 +677,67 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_common_stream(
|
||||
streams: &mut tokio::task::JoinSet<Result<(), StreamableHttpError<C::Error>>>,
|
||||
client: C,
|
||||
session_id: Arc<str>,
|
||||
config: &StreamableHttpClientTransportConfig,
|
||||
protocol_headers: HashMap<HeaderName, HeaderValue>,
|
||||
sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>,
|
||||
transport_task_ct: CancellationToken,
|
||||
) {
|
||||
let uri = config.uri.clone();
|
||||
let auth_header = config.auth_header.clone();
|
||||
let retry_config = config.retry_config.clone();
|
||||
let reconnect_uri = config.uri.clone();
|
||||
let reconnect_auth_header = config.auth_header.clone();
|
||||
let max_sse_event_size = config.max_sse_event_size;
|
||||
|
||||
streams.spawn(async move {
|
||||
match client
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
session_id.clone(),
|
||||
None,
|
||||
auth_header,
|
||||
protocol_headers.clone(),
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(stream) => {
|
||||
let sse_stream = SseAutoReconnectStream::new(
|
||||
stream,
|
||||
StreamableHttpClientReconnect {
|
||||
client,
|
||||
session_id,
|
||||
uri: reconnect_uri,
|
||||
auth_header: reconnect_auth_header,
|
||||
custom_headers: protocol_headers,
|
||||
max_sse_event_size,
|
||||
},
|
||||
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(error) => {
|
||||
tracing::error!("fail to get common stream: {error}");
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Performs a transparent re-initialization handshake after a session-expired 404.
|
||||
///
|
||||
/// Takes an owned clone of the client (avoiding `&self` across `.await` so the
|
||||
|
|
@ -667,8 +754,14 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
auth_header: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<(Option<Arc<str>>, HashMap<HeaderName, HeaderValue>), StreamableHttpError<C::Error>>
|
||||
{
|
||||
) -> Result<
|
||||
(
|
||||
Option<Arc<str>>,
|
||||
ProtocolVersion,
|
||||
HashMap<HeaderName, HeaderValue>,
|
||||
),
|
||||
StreamableHttpError<C::Error>,
|
||||
> {
|
||||
let (init_msg, new_session_id_str) = client
|
||||
.post_message_with_max_sse_event_size(
|
||||
uri.clone(),
|
||||
|
|
@ -712,7 +805,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
.await?
|
||||
.expect_accepted_or_json::<C::Error>()?;
|
||||
|
||||
Ok((new_session_id, new_protocol_headers))
|
||||
Ok((new_session_id, negotiated_version, new_protocol_headers))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -743,17 +836,33 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
let _drop_guard = transport_task_ct.clone().drop_guard();
|
||||
let WorkerSendRequest {
|
||||
responder,
|
||||
message: initialize_request,
|
||||
message: startup_request,
|
||||
} = context.recv_from_handler().await?;
|
||||
let saved_init_request = initialize_request.clone();
|
||||
let is_legacy_startup = matches!(
|
||||
&startup_request,
|
||||
ClientJsonRpcMessage::Request(request)
|
||||
if matches!(&request.request, ClientRequest::InitializeRequest(_))
|
||||
);
|
||||
let mut saved_init_request = is_legacy_startup.then(|| startup_request.clone());
|
||||
let empty_tool_cache = HashMap::new();
|
||||
let (bootstrap_version, bootstrap_headers) = if is_legacy_startup {
|
||||
(ProtocolVersion::default(), config.custom_headers.clone())
|
||||
} else {
|
||||
request_version_headers(
|
||||
&config.custom_headers,
|
||||
&startup_request,
|
||||
&ProtocolVersion::default(),
|
||||
&empty_tool_cache,
|
||||
)
|
||||
};
|
||||
let (message, session_id) = match self
|
||||
.client
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
initialize_request,
|
||||
startup_request,
|
||||
None,
|
||||
config.auth_header.clone(),
|
||||
config.custom_headers.clone(),
|
||||
bootstrap_headers.clone(),
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
|
|
@ -784,8 +893,11 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
}
|
||||
None
|
||||
};
|
||||
let (negotiated_version, mut protocol_headers) =
|
||||
negotiate_version_headers(&message, config.custom_headers.clone());
|
||||
let (mut negotiated_version, mut protocol_headers) = if is_legacy_startup {
|
||||
negotiate_version_headers(&message, config.custom_headers.clone())
|
||||
} else {
|
||||
(bootstrap_version, bootstrap_headers)
|
||||
};
|
||||
// SEP-2243: tool input schemas (name -> schema) cached from tools/list responses,
|
||||
// used to promote annotated tools/call arguments to Mcp-Param-* headers.
|
||||
let mut tool_header_cache: HashMap<String, Arc<JsonObject>> = HashMap::new();
|
||||
|
|
@ -800,33 +912,37 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
});
|
||||
|
||||
context.send_to_handler(message).await?;
|
||||
let initialized_notification = context.recv_from_handler().await?;
|
||||
// expect a initialized response
|
||||
let initialized_headers = build_request_headers(
|
||||
&protocol_headers,
|
||||
&initialized_notification.message,
|
||||
&tool_header_cache,
|
||||
&negotiated_version,
|
||||
);
|
||||
self.client
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
initialized_notification.message,
|
||||
session_id.clone(),
|
||||
config.auth_header.clone(),
|
||||
initialized_headers,
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
.map_err(WorkerQuitReason::fatal_context(
|
||||
"send initialized notification",
|
||||
))?
|
||||
.expect_accepted_or_json::<C::Error>()
|
||||
.map_err(WorkerQuitReason::fatal_context(
|
||||
"process initialized notification response",
|
||||
))?;
|
||||
let _ = initialized_notification.responder.send(Ok(()));
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
if is_legacy_startup {
|
||||
let initialized_notification = context.recv_from_handler().await?;
|
||||
let initialized_headers = build_request_headers(
|
||||
&protocol_headers,
|
||||
&initialized_notification.message,
|
||||
&tool_header_cache,
|
||||
&negotiated_version,
|
||||
);
|
||||
self.client
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
initialized_notification.message,
|
||||
session_id.clone(),
|
||||
config.auth_header.clone(),
|
||||
initialized_headers,
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
.map_err(WorkerQuitReason::fatal_context(
|
||||
"send initialized notification",
|
||||
))?
|
||||
.expect_accepted_or_json::<C::Error>()
|
||||
.map_err(WorkerQuitReason::fatal_context(
|
||||
"process initialized notification response",
|
||||
))?;
|
||||
let _ = initialized_notification.responder.send(Ok(()));
|
||||
}
|
||||
#[expect(
|
||||
clippy::large_enum_variant,
|
||||
reason = "the event is short-lived and boxing would add allocation in the event loop"
|
||||
)]
|
||||
enum Event<W: Worker, E: std::error::Error + Send + Sync + 'static> {
|
||||
ClientMessage(WorkerSendRequest<W>),
|
||||
ServerMessage(ServerJsonRpcMessage),
|
||||
|
|
@ -834,63 +950,17 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
}
|
||||
let mut streams = tokio::task::JoinSet::new();
|
||||
let mut pending_stream_response_ids = HashSet::new();
|
||||
let mut awaiting_fallback_initialized = false;
|
||||
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();
|
||||
let spawn_headers = protocol_headers.clone();
|
||||
let max_sse_event_size = config.max_sse_event_size;
|
||||
|
||||
streams.spawn(async move {
|
||||
match client
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri.clone(),
|
||||
session_id.clone(),
|
||||
None,
|
||||
auth_header.clone(),
|
||||
spawn_headers.clone(),
|
||||
max_sse_event_size,
|
||||
)
|
||||
.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,
|
||||
custom_headers: spawn_headers,
|
||||
max_sse_event_size,
|
||||
},
|
||||
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)
|
||||
}
|
||||
}
|
||||
});
|
||||
Self::spawn_common_stream(
|
||||
&mut streams,
|
||||
self.client.clone(),
|
||||
session_id.clone(),
|
||||
&config,
|
||||
protocol_headers.clone(),
|
||||
sse_worker_tx.clone(),
|
||||
transport_task_ct.clone(),
|
||||
);
|
||||
}
|
||||
// Main event loop - capture exit reason so we can do cleanup before returning
|
||||
let loop_result: Result<(), WorkerQuitReason<Self::Error>> = 'main_loop: loop {
|
||||
|
|
@ -926,16 +996,112 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
match event {
|
||||
Event::ClientMessage(send_request) => {
|
||||
let WorkerSendRequest { message, responder } = send_request;
|
||||
let is_fallback_initialize = saved_init_request.is_none()
|
||||
&& matches!(
|
||||
&message,
|
||||
ClientJsonRpcMessage::Request(request)
|
||||
if matches!(
|
||||
&request.request,
|
||||
ClientRequest::InitializeRequest(_)
|
||||
)
|
||||
);
|
||||
if is_fallback_initialize {
|
||||
saved_init_request = Some(message.clone());
|
||||
// Servers do not assign sessions to `server/discover`, so a
|
||||
// fallback initialize starts from a clean slate: no session
|
||||
// ID, no cleanup state, and no streams to tear down.
|
||||
debug_assert!(
|
||||
session_id.is_none()
|
||||
&& session_cleanup_info.is_none()
|
||||
&& streams.is_empty(),
|
||||
"discover bootstrap must not create session state"
|
||||
);
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
message,
|
||||
None,
|
||||
config.auth_header.clone(),
|
||||
config.custom_headers.clone(),
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await;
|
||||
let response = match response {
|
||||
Ok(response) => {
|
||||
let _ = responder.send(Ok(()));
|
||||
response
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = responder.send(Err(error));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let (initialize_response, new_session_id) = response
|
||||
.expect_initialized::<C::Error>()
|
||||
.await
|
||||
.map_err(WorkerQuitReason::fatal_context(
|
||||
"process fallback initialize response",
|
||||
))?;
|
||||
session_id = new_session_id.map(Arc::from);
|
||||
if session_id.is_none() && !config.allow_stateless {
|
||||
return Err(WorkerQuitReason::fatal(
|
||||
StreamableHttpError::<C::Error>::MissingSessionIdInResponse,
|
||||
"process fallback initialize response",
|
||||
));
|
||||
}
|
||||
(negotiated_version, protocol_headers) = negotiate_version_headers(
|
||||
&initialize_response,
|
||||
config.custom_headers.clone(),
|
||||
);
|
||||
session_cleanup_info =
|
||||
session_id.as_ref().map(|session_id| SessionCleanupInfo {
|
||||
client: self.client.clone(),
|
||||
uri: config.uri.clone(),
|
||||
session_id: session_id.clone(),
|
||||
auth_header: config.auth_header.clone(),
|
||||
protocol_headers: protocol_headers.clone(),
|
||||
});
|
||||
context.send_to_handler(initialize_response).await?;
|
||||
awaiting_fallback_initialized = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
let request_id = Self::client_request_id(&message);
|
||||
let inline_version = match &message {
|
||||
ClientJsonRpcMessage::Request(request) => {
|
||||
request.request.get_meta().protocol_version()
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let is_initialized_notification = matches!(
|
||||
&message,
|
||||
ClientJsonRpcMessage::Notification(notification)
|
||||
if matches!(
|
||||
¬ification.notification,
|
||||
ClientNotification::InitializedNotification(_)
|
||||
)
|
||||
);
|
||||
// Pass a clone to the first attempt so `message` is retained for a
|
||||
// potential re-init retry. `post_message` takes ownership and the
|
||||
// trait cannot be changed, so the clone is unavoidable.
|
||||
let request_headers = build_request_headers(
|
||||
let (request_version, request_headers) = request_version_headers(
|
||||
&protocol_headers,
|
||||
&message,
|
||||
&tool_header_cache,
|
||||
&negotiated_version,
|
||||
&tool_header_cache,
|
||||
);
|
||||
if inline_version.is_some() {
|
||||
negotiated_version = request_version.clone();
|
||||
if let Ok(value) = HeaderValue::from_str(request_version.as_str()) {
|
||||
protocol_headers
|
||||
.insert(HeaderName::from_static("mcp-protocol-version"), value);
|
||||
}
|
||||
if let Some(cleanup) = &mut session_cleanup_info {
|
||||
cleanup.protocol_headers = protocol_headers.clone();
|
||||
}
|
||||
}
|
||||
let response = self
|
||||
.client
|
||||
.post_message_with_max_sse_event_size(
|
||||
|
|
@ -949,9 +1115,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
.await;
|
||||
let send_result = match response {
|
||||
Err(StreamableHttpError::SessionExpired) => {
|
||||
if !config.reinit_on_expired_session {
|
||||
Err(StreamableHttpError::SessionExpired)
|
||||
} else {
|
||||
if let Some(saved_init_request) = saved_init_request
|
||||
.as_ref()
|
||||
.filter(|_| config.reinit_on_expired_session)
|
||||
{
|
||||
// The server discarded the session (HTTP 404). Perform a
|
||||
// fresh handshake once and replay the original message.
|
||||
tracing::info!(
|
||||
|
|
@ -967,7 +1134,11 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
)
|
||||
.await
|
||||
{
|
||||
Ok((new_session_id, new_protocol_headers)) => {
|
||||
Ok((
|
||||
new_session_id,
|
||||
new_negotiated_version,
|
||||
new_protocol_headers,
|
||||
)) => {
|
||||
// Old streams hold the stale session ID. Stop them first
|
||||
// so no late stale-session messages can arrive after the
|
||||
// pending requests below are completed.
|
||||
|
|
@ -990,6 +1161,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
.await?;
|
||||
|
||||
session_id = new_session_id;
|
||||
negotiated_version = new_negotiated_version;
|
||||
protocol_headers = new_protocol_headers;
|
||||
session_cleanup_info =
|
||||
session_id.as_ref().map(|sid| SessionCleanupInfo {
|
||||
|
|
@ -1001,71 +1173,22 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
});
|
||||
|
||||
if let Some(new_sid) = &session_id {
|
||||
let client = self.client.clone();
|
||||
let uri = config.uri.clone();
|
||||
let new_sid = new_sid.clone();
|
||||
let auth_header = config.auth_header.clone();
|
||||
let retry_config = self.config.retry_config.clone();
|
||||
let sse_tx = sse_worker_tx.clone();
|
||||
let task_ct = transport_task_ct.clone();
|
||||
let config_uri = config.uri.clone();
|
||||
let config_auth = config.auth_header.clone();
|
||||
let spawn_headers = protocol_headers.clone();
|
||||
let max_sse_event_size = config.max_sse_event_size;
|
||||
streams.spawn(async move {
|
||||
match client
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
new_sid.clone(),
|
||||
None,
|
||||
auth_header.clone(),
|
||||
spawn_headers.clone(),
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(stream) => {
|
||||
let sse_stream = SseAutoReconnectStream::new(
|
||||
stream,
|
||||
StreamableHttpClientReconnect {
|
||||
client: client.clone(),
|
||||
session_id: new_sid,
|
||||
uri: config_uri,
|
||||
auth_header: config_auth,
|
||||
custom_headers: spawn_headers,
|
||||
max_sse_event_size,
|
||||
},
|
||||
retry_config,
|
||||
);
|
||||
Self::execute_sse_stream(
|
||||
sse_stream,
|
||||
sse_tx,
|
||||
false,
|
||||
task_ct.child_token(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
Err(StreamableHttpError::ServerDoesNotSupportSse) => {
|
||||
tracing::debug!(
|
||||
"server doesn't support sse after re-init"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
"fail to get common stream after re-init: {e}"
|
||||
);
|
||||
Err(e)
|
||||
}
|
||||
}
|
||||
});
|
||||
Self::spawn_common_stream(
|
||||
&mut streams,
|
||||
self.client.clone(),
|
||||
new_sid.clone(),
|
||||
&config,
|
||||
protocol_headers.clone(),
|
||||
sse_worker_tx.clone(),
|
||||
transport_task_ct.clone(),
|
||||
);
|
||||
}
|
||||
|
||||
let retry_headers = build_request_headers(
|
||||
let (_, retry_headers) = request_version_headers(
|
||||
&protocol_headers,
|
||||
&message,
|
||||
&tool_header_cache,
|
||||
&negotiated_version,
|
||||
&tool_header_cache,
|
||||
);
|
||||
let retry_response = self
|
||||
.client
|
||||
|
|
@ -1126,7 +1249,9 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
}
|
||||
Err(reinit_err) => Err(reinit_err),
|
||||
}
|
||||
} // else enable_reinit_on_expired_session
|
||||
} else {
|
||||
Err(StreamableHttpError::SessionExpired)
|
||||
}
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
Ok(StreamableHttpPostResponse::Accepted) => {
|
||||
|
|
@ -1167,6 +1292,23 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
Ok(())
|
||||
}
|
||||
};
|
||||
if send_result.is_ok()
|
||||
&& awaiting_fallback_initialized
|
||||
&& is_initialized_notification
|
||||
{
|
||||
if let Some(session_id) = &session_id {
|
||||
Self::spawn_common_stream(
|
||||
&mut streams,
|
||||
self.client.clone(),
|
||||
session_id.clone(),
|
||||
&config,
|
||||
protocol_headers.clone(),
|
||||
sse_worker_tx.clone(),
|
||||
transport_task_ct.clone(),
|
||||
);
|
||||
}
|
||||
awaiting_fallback_initialized = false;
|
||||
}
|
||||
let _ = responder.send(send_result);
|
||||
}
|
||||
Event::ServerMessage(json_rpc_message) => {
|
||||
|
|
@ -1194,7 +1336,6 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
// 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(cleanup) = session_cleanup_info {
|
||||
const SESSION_CLEANUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
|
||||
let cleanup_session_id = cleanup.session_id.clone();
|
||||
match tokio::time::timeout(
|
||||
SESSION_CLEANUP_TIMEOUT,
|
||||
|
|
|
|||
Loading…
Reference in a new issue