1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
//! Routes MCP `notifications/progress` into the SDK tool-progress channel.
//!
//! rmcp stamps every outbound request with a progress token, so a server can
//! report progress against any in-flight call. The router maps the token of a
//! live `tools/call` to that invocation's [`ToolProgressSender`], which is what
//! renders the live line on a tool card. Subscriptions are scoped to the call
//! through [`ProgressSubscription`], so a finished or cancelled call always
//! stops receiving.
use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
use rho_sdk::tool::{ToolProgress, ToolProgressSender};
use rmcp::model::{ProgressNotificationParam, ProgressToken};
/// Live progress subscriptions for one MCP session.
#[derive(Clone, Debug, Default)]
pub(crate) struct McpProgressRouter {
subscribers: Arc<Mutex<HashMap<ProgressToken, ToolProgressSender>>>,
}
impl McpProgressRouter {
pub(crate) fn new() -> Self {
Self::default()
}
/// Route progress for `token` to `sender` until the guard drops.
pub(crate) fn subscribe(
&self,
token: ProgressToken,
sender: ToolProgressSender,
) -> ProgressSubscription {
self.lock().insert(token.clone(), sender);
ProgressSubscription {
token,
router: self.clone(),
}
}
pub(crate) async fn dispatch(&self, params: ProgressNotificationParam) {
let Some(sender) = self.lock().get(¶ms.progress_token).cloned() else {
return;
};
// The lock is released before the await: `send` applies backpressure and
// must never hold the map across a suspension point.
sender.send(tool_progress(params)).await;
}
fn unsubscribe(&self, token: &ProgressToken) {
self.lock().remove(token);
}
/// A poisoned lock means a panic while a subscription map was borrowed. The
/// map stays structurally valid, so recover rather than fail a tool call.
fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<ProgressToken, ToolProgressSender>> {
self.subscribers
.lock()
.unwrap_or_else(|error| error.into_inner())
}
}
/// Drops the routing entry when one tool call ends.
pub(crate) struct ProgressSubscription {
token: ProgressToken,
router: McpProgressRouter,
}
impl Drop for ProgressSubscription {
fn drop(&mut self) {
self.router.unsubscribe(&self.token);
}
}
/// MCP progress is a float pair; the SDK card shows whole units. Counts are
/// only forwarded when the server supplies a total, because a bare growing
/// float has nothing to render against.
fn tool_progress(params: ProgressNotificationParam) -> ToolProgress {
let message = params
.message
.unwrap_or_else(|| "MCP server reported progress".into());
let progress = ToolProgress::message(message);
match params.total {
Some(total) if total.is_finite() && total > 0.0 => {
progress.units(whole_units(params.progress), whole_units(total))
}
_ => progress,
}
}
fn whole_units(value: f64) -> u64 {
if !value.is_finite() || value <= 0.0 {
return 0;
}
value.round() as u64
}
#[cfg(test)]
#[path = "progress_tests.rs"]
mod tests;