1use super::McpClient;
2use super::task::{TaskErrorReason, cancel_server_task, stream_task_events};
3use async_stream::stream;
4use futures::future::Either;
5use futures::{Stream, StreamExt, stream};
6use rmcp::RoleClient;
7use rmcp::model::{
8 CallToolRequestParams, CallToolResponse, CallToolResult, ClientRequest, CreateTaskResult, DEFAULT_MRTR_MAX_ROUNDS,
9 InputRequests, InputRequiredResult, ProgressNotificationParam, Request, RequestMetaObject, ServerResult, Task,
10};
11use rmcp::service::{PeerRequestOptions, RequestHandle, ServiceError};
12use std::future::pending;
13use std::pin::{Pin, pin};
14use std::task::{Context, Poll};
15use std::time::Duration;
16use thiserror::Error;
17use tokio::time::{Instant, sleep, sleep_until};
18use tokio_util::sync::CancellationToken;
19
20pub struct ToolCall {
21 events: Pin<Box<dyn Stream<Item = ToolCallEvent> + Send>>,
22}
23
24#[derive(Debug)]
25pub enum ToolCallEvent {
26 Progress(ProgressNotificationParam),
27 TaskCreated(CreateTaskResult),
28 TaskStatus(Task),
29 Done(Result<CallToolResult, ToolCallError>),
30}
31
32#[derive(Debug, Clone, Default)]
33pub struct ToolCallOptions {
34 pub timeout: Option<Duration>,
35 pub meta: Option<RequestMetaObject>,
36 pub cancel: CancellationToken,
37}
38
39#[derive(Debug, Error)]
40pub enum ToolCallError {
41 #[error("Failed to send tool request: {0}")]
42 Send(#[source] ServiceError),
43 #[error("Tool execution failed: {0}")]
44 Call(#[source] ServiceError),
45 #[error("Tool call exceeded its {0:?} timeout")]
46 TimedOut(Duration),
47 #[error("Server requested an input kind this client does not support (sampling or roots)")]
48 UnsupportedInput,
49 #[error("Server requested input without any input requests or request state")]
50 EmptyInputRequired,
51 #[error("Server requested input again after the user cancelled")]
52 RePromptAfterCancel,
53 #[error("Server did not complete within {DEFAULT_MRTR_MAX_ROUNDS} MRTR input rounds")]
54 InputRoundsExceeded,
55 #[error("Task '{task_id}' {reason}")]
56 Task {
57 task_id: String,
58 #[source]
59 reason: Box<TaskErrorReason>,
60 },
61 #[error("Server returned a tool call response kind this client does not support")]
62 UnsupportedResponse,
63 #[error("Tool call was cancelled")]
64 Cancelled,
65}
66
67impl ToolCallOptions {
68 pub fn with_timeout(timeout: Duration) -> Self {
69 Self { timeout: Some(timeout), ..Self::default() }
70 }
71}
72
73impl ToolCall {
74 pub(crate) fn new(client: McpClient, params: CallToolRequestParams, options: ToolCallOptions) -> Self {
75 let ToolCallOptions { timeout, meta, cancel } = options;
76 let events = stream_events(client.clone(), params, meta);
77 Self { events: Box::pin(bounded(client, events, timeout, cancel)) }
78 }
79
80 pub async fn result(self) -> Result<CallToolResult, ToolCallError> {
81 self.result_with_progress(|_| std::future::ready(())).await
82 }
83
84 pub async fn result_with_progress<F: Future<Output = ()>>(
85 mut self,
86 mut on_progress: impl FnMut(ProgressNotificationParam) -> F,
87 ) -> Result<CallToolResult, ToolCallError> {
88 while let Some(event) = self.next().await {
89 match event {
90 ToolCallEvent::Progress(progress) => on_progress(progress).await,
91 ToolCallEvent::Done(result) => return result,
92 ToolCallEvent::TaskCreated(_) | ToolCallEvent::TaskStatus(_) => {}
93 }
94 }
95 Err(ToolCallError::Cancelled)
96 }
97}
98
99impl Stream for ToolCall {
100 type Item = ToolCallEvent;
101
102 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
103 self.events.as_mut().poll_next(cx)
104 }
105}
106
107impl From<ToolCallError> for ToolCallEvent {
108 fn from(error: ToolCallError) -> Self {
109 Self::Done(Err(error))
110 }
111}
112
113pub(super) fn with_progress<'a, T: 'a>(
114 progress: impl Stream<Item = ProgressNotificationParam> + 'a,
115 events: impl Stream<Item = T> + 'a,
116) -> impl Stream<Item = Either<ProgressNotificationParam, T>> + 'a {
117 stream! {
118 let mut progress = pin!(progress);
119 let mut events = pin!(events);
120 loop {
121 tokio::select! {
122 biased;
123 progress_event = progress.next() => match progress_event {
124 Some(progress_event) => yield Either::Left(progress_event),
125 None => {
126 while let Some(event) = events.next().await {
127 yield Either::Right(event);
128 }
129 return;
130 }
131 },
132 event = events.next() => match event {
133 Some(event) => yield Either::Right(event),
134 None => return,
135 },
136 }
137 }
138 }
139}
140
141fn bounded(
142 client: McpClient,
143 events: impl Stream<Item = ToolCallEvent> + Send + 'static,
144 timeout: Option<Duration>,
145 cancel: CancellationToken,
146) -> impl Stream<Item = ToolCallEvent> + Send + 'static {
147 stream! {
148 let mut events = Box::pin(events);
149 let mut expired = pin!(expiry(timeout));
150 let mut cancelled = pin!(cancel.cancelled());
151 let mut task = None;
152
153 let interrupted = loop {
154 tokio::select! {
155 biased;
156 event = events.next() => {
157 let Some(event) = event else { return };
158 match &event {
159 ToolCallEvent::TaskCreated(created) => task = Some(created.task.clone()),
160 ToolCallEvent::TaskStatus(status) => task = Some(status.clone()),
161 ToolCallEvent::Progress(_) | ToolCallEvent::Done(_) => {}
162 }
163 let done = matches!(event, ToolCallEvent::Done(_));
164 yield event;
165 if done {
166 return;
167 }
168 }
169 timeout = &mut expired => break ToolCallError::TimedOut(timeout),
170 () = &mut cancelled => break ToolCallError::Cancelled,
171 }
172 };
173
174 drop(events);
175 if let Some(task) = &task {
176 cancel_server_task(&client, &task.task_id).await;
177 }
178 yield interrupted.into();
179 }
180}
181
182fn stream_events(
183 client: McpClient,
184 mut params: CallToolRequestParams,
185 meta: Option<RequestMetaObject>,
186) -> impl Stream<Item = ToolCallEvent> + Send + 'static {
187 stream! {
188 let mut mrtr_state = MrtrState::new();
189
190 loop {
191 let request = ClientRequest::CallToolRequest(Request::new(params.clone()));
192 let handle = match client.peer.send_cancellable_request(request, peer_request_options(meta.clone())).await {
193 Ok(handle) => handle,
194 Err(error) => {
195 yield ToolCallError::Send(error).into();
196 return;
197 }
198 };
199
200 let mut progress = client.handler.progress.subscribe(handle.progress_token.clone()).await;
201 let mut response = None;
202 {
203 let mut awaited = pin!(with_progress(&mut progress, stream::once(await_tool_response(handle))));
204 while let Some(event) = awaited.next().await {
205 match event {
206 Either::Left(progress_event) => yield ToolCallEvent::Progress(progress_event),
207 Either::Right(awaited_response) => response = Some(awaited_response),
208 }
209 }
210 }
211
212 match response.unwrap_or(Err(ServiceError::UnexpectedResponse)) {
213 Ok(CallToolResponse::Complete(result)) => {
214 yield ToolCallEvent::Done(Ok(result));
215 return;
216 }
217 Ok(CallToolResponse::InputRequired(input_required)) => match mrtr_state.tick(input_required) {
218 MrtrAction::Poll { backoff, request_state } => {
219 sleep(backoff).await;
220 params.input_responses = None;
221 params.request_state = Some(request_state);
222 }
223 MrtrAction::Elicit { input_requests, request_state } => {
224 let (responses, cancelled) = match client.handler.elicit_inputs(input_requests).await {
225 Ok(answered) => answered,
226 Err(error) => {
227 yield error.into();
228 return;
229 }
230 };
231 mrtr_state.record_cancelled(cancelled);
232 params.input_responses = Some(responses);
233 params.request_state = request_state;
234 }
235 MrtrAction::Abort(error) => {
236 yield error.into();
237 return;
238 }
239 },
240 Ok(CallToolResponse::Task(created)) => {
241 let mut events = pin!(stream_task_events(&client, created, progress));
242 while let Some(event) = events.next().await {
243 yield event;
244 }
245 return;
246 }
247 Ok(_) => {
248 yield ToolCallError::UnsupportedResponse.into();
249 return;
250 }
251 Err(error) => {
252 yield ToolCallError::Call(error).into();
253 return;
254 }
255 }
256 }
257 }
258}
259
260async fn expiry(timeout: Option<Duration>) -> Duration {
261 match timeout.and_then(|timeout| Some((Instant::now().checked_add(timeout)?, timeout))) {
262 Some((deadline, timeout)) => {
263 sleep_until(deadline).await;
264 timeout
265 }
266 None => pending().await,
267 }
268}
269
270fn peer_request_options(meta: Option<RequestMetaObject>) -> PeerRequestOptions {
271 meta.map_or_else(PeerRequestOptions::no_options, |meta| PeerRequestOptions::no_options().with_meta(meta))
272}
273
274async fn await_tool_response(handle: RequestHandle<RoleClient>) -> Result<CallToolResponse, ServiceError> {
275 match handle.await_response().await? {
276 ServerResult::CallToolResult(result) => Ok(CallToolResponse::Complete(result)),
277 ServerResult::InputRequiredResult(result) => Ok(CallToolResponse::InputRequired(result)),
278 ServerResult::CreateTaskResult(result) => Ok(CallToolResponse::Task(result)),
279 _ => Err(ServiceError::UnexpectedResponse),
280 }
281}
282
283struct MrtrState {
284 input_request_rounds: usize,
285 next_backoff: Duration,
286 user_cancelled: bool,
287}
288
289enum MrtrAction {
290 Poll { backoff: Duration, request_state: String },
291 Elicit { input_requests: InputRequests, request_state: Option<String> },
292 Abort(ToolCallError),
293}
294
295impl MrtrState {
296 fn new() -> Self {
297 Self { input_request_rounds: 0, next_backoff: BASE_BACKOFF, user_cancelled: false }
298 }
299
300 fn tick(&mut self, input_required: InputRequiredResult) -> MrtrAction {
301 let input_requests = input_required.input_requests.filter(|requests| !requests.is_empty());
302 match (input_requests, input_required.request_state) {
303 (None, None) => MrtrAction::Abort(ToolCallError::EmptyInputRequired),
304 (None, Some(request_state)) => {
305 let backoff = self.next_backoff;
306 self.next_backoff = (backoff * 2).min(MAX_BACKOFF);
307 MrtrAction::Poll { backoff, request_state }
308 }
309 (Some(input_requests), request_state) => {
310 if self.user_cancelled {
311 MrtrAction::Abort(ToolCallError::RePromptAfterCancel)
312 } else if self.input_request_rounds == DEFAULT_MRTR_MAX_ROUNDS {
313 MrtrAction::Abort(ToolCallError::InputRoundsExceeded)
314 } else {
315 self.input_request_rounds += 1;
316 self.next_backoff = BASE_BACKOFF;
317 MrtrAction::Elicit { input_requests, request_state }
318 }
319 }
320 }
321 }
322
323 fn record_cancelled(&mut self, cancelled: bool) {
324 self.user_cancelled |= cancelled;
325 }
326}
327
328const BASE_BACKOFF: Duration = Duration::from_millis(50);
329const MAX_BACKOFF: Duration = Duration::from_millis(1600);