portugalfuturista/rust-sdk
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: align progress timeout token (#909)

Browse source
This commit is contained in:
Dale Seo 2026-06-22 10:59:32 -04:00 committed by GitHub
parent 4fd4986b62
commit 443677ca31
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 42 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
crates/rmcp/src/service.rs
View file

@ -348,7 +348,6 @@ pub struct RequestHandle<R: ServiceRole> {
pub peer: Peer<R>,
pub id: RequestId,
pub progress_token: ProgressToken,
progress_timeout_watchers: ProgressTimeoutWatchers,
progress_reset_rx: Option<mpsc::Receiver<()>>,
}
@ -416,8 +415,10 @@ impl<R: ServiceRole> RequestHandle<R> {
max_total_timeout: Option<Duration>,
reset_timeout_on_progress: bool,
) -> Result<R::PeerResp, ServiceError> {
let mut idle_sleep = timeout.map(tokio::time::sleep).map(Box::pin);
let mut max_total_sleep = max_total_timeout.map(tokio::time::sleep).map(Box::pin);
let mut idle_sleep =
timeout.map(|timeout| (timeout, Box::pin(tokio::time::sleep(timeout))));
let mut max_total_sleep =
max_total_timeout.map(|timeout| (timeout, Box::pin(tokio::time::sleep(timeout))));
loop {
tokio::select! {
@ -427,32 +428,34 @@ impl<R: ServiceRole> RequestHandle<R> {
return response.map_err(|_e| ServiceError::TransportClosed)?;
}
_ = async {
if let Some(sleep) = idle_sleep.as_mut() {
if let Some((_, sleep)) = idle_sleep.as_mut() {
sleep.as_mut().await;
}
}, if idle_sleep.is_some() => {
let timeout = timeout.expect("idle timeout exists when idle sleep exists");
self.send_timeout_cancel_notification(Self::REQUEST_TIMEOUT_REASON).await;
return Err(ServiceError::Timeout { timeout });
if let Some((timeout, _)) = idle_sleep.as_ref() {
self.send_timeout_cancel_notification(Self::REQUEST_TIMEOUT_REASON).await;
return Err(ServiceError::Timeout { timeout: *timeout });
}
}
_ = async {
if let Some(sleep) = max_total_sleep.as_mut() {
if let Some((_, sleep)) = max_total_sleep.as_mut() {
sleep.as_mut().await;
}
}, if max_total_sleep.is_some() => {
let timeout = max_total_timeout.expect("max total timeout exists when max total sleep exists");
self.send_timeout_cancel_notification(Self::REQUEST_MAX_TOTAL_TIMEOUT_REASON).await;
return Err(ServiceError::Timeout { timeout });
if let Some((timeout, _)) = max_total_sleep.as_ref() {
self.send_timeout_cancel_notification(Self::REQUEST_MAX_TOTAL_TIMEOUT_REASON).await;
return Err(ServiceError::Timeout { timeout: *timeout });
}
}
progress = async {
match self.progress_reset_rx.as_mut() {
Some(rx) => rx.recv().await,
None => None,
}
}, if reset_timeout_on_progress && timeout.is_some() && self.progress_reset_rx.is_some() => {
}, if reset_timeout_on_progress && idle_sleep.is_some() && self.progress_reset_rx.is_some() => {
if progress.is_some() {
if let (Some(timeout), Some(sleep)) = (timeout, idle_sleep.as_mut()) {
sleep.as_mut().reset(tokio::time::Instant::now() + timeout);
if let Some((timeout, sleep)) = idle_sleep.as_mut() {
sleep.as_mut().reset(tokio::time::Instant::now() + *timeout);
}
}
}
@ -463,7 +466,7 @@ impl<R: ServiceRole> RequestHandle<R> {
/// Cancel this request
pub async fn cancel(self, reason: Option<String>) -> Result<(), ServiceError> {
Self::cleanup_progress_timeout_watcher(
&self.progress_timeout_watchers,
&self.peer.progress_timeout_watchers,
&self.progress_token,
self.progress_reset_rx.is_some(),
)
@ -617,12 +620,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();
request
.get_meta_mut()
.set_progress_token(progress_token.clone());
if let Some(meta) = options.meta.clone() {
request.get_meta_mut().extend(meta);
}
request
.get_meta_mut()
.set_progress_token(progress_token.clone());
let (responder, receiver) = tokio::sync::oneshot::channel();
let progress_reset_rx = if options.reset_timeout_on_progress && options.timeout.is_some() {
let (sender, receiver) = mpsc::channel(1);
@ -658,7 +661,6 @@ impl<R: ServiceRole> Peer<R> {
progress_token,
options,
peer: self.clone(),
progress_timeout_watchers: self.progress_timeout_watchers.clone(),
progress_reset_rx,
})
}

32
crates/rmcp/tests/test_request_timeout_progress.rs
View file

@ -10,7 +10,10 @@ use std::{
use rmcp::{
ClientHandler, Peer, RoleServer, ServiceError, ServiceExt,
model::{CallToolRequestParams, ClientRequest, Meta, ProgressNotificationParam, Request},
model::{
CallToolRequestParams, ClientRequest, Meta, NumberOrString, ProgressNotificationParam,
ProgressToken, Request,
},
service::PeerRequestOptions,
tool, tool_router,
};
@ -32,12 +35,6 @@ impl ClientHandler for ProgressCountingClient {
struct ProgressTimeoutServer;
impl ProgressTimeoutServer {
fn new() -> Self {
Self
}
}
#[tool_router(server_handler)]
impl ProgressTimeoutServer {
#[tool]
@ -83,9 +80,7 @@ impl ProgressTimeoutServer {
tokio::time::sleep(Duration::from_millis(50)).await;
let _ = client
.notify_progress(ProgressNotificationParam {
progress_token: rmcp::model::ProgressToken(
rmcp::model::NumberOrString::Number(999_999),
),
progress_token: ProgressToken(NumberOrString::Number(999_999)),
progress: step as f64,
total: Some(4.0),
message: Some("unrelated".into()),
@ -99,7 +94,7 @@ impl ProgressTimeoutServer {
async fn start_pair()
-> anyhow::Result<rmcp::service::RunningService<rmcp::RoleClient, ProgressCountingClient>> {
let server = ProgressTimeoutServer::new();
let server = ProgressTimeoutServer;
let client = ProgressCountingClient::default();
let (transport_server, transport_client) = tokio::io::duplex(4096);
@ -172,6 +167,21 @@ async fn matching_progress_resets_timeout_when_enabled() -> anyhow::Result<()> {
Ok(())
}
#[tokio::test]
async fn generated_progress_token_overrides_option_meta_token() -> anyhow::Result<()> {
let client = start_pair().await?;
let mut options =
PeerRequestOptions::with_timeout(Duration::from_millis(75)).reset_timeout_on_progress();
options.meta = Some(Meta::with_progress_token(ProgressToken(
NumberOrString::Number(999_999),
)));
let result = call_tool_with_options(&client, "delayed_with_progress", options).await;
assert!(result.is_ok());
Ok(())
}
#[tokio::test]
async fn max_total_timeout_wins_over_progress_reset() -> anyhow::Result<()> {
let client = start_pair().await?;