mcp_utils/client/
task.rs

1use super::McpClient;
2use super::tool_call::{ToolCallError, ToolCallEvent, with_progress};
3use async_stream::stream;
4use futures::future::Either;
5use futures::{Stream, StreamExt};
6use rmcp::model::{
7    CancelTaskParams, CreateTaskResult, GetTaskParams, InputRequests, ProgressNotificationParam, Task, TaskPayload,
8    TaskStatus, UpdateTaskParams,
9};
10use rmcp::service::ServiceError;
11use std::collections::HashSet;
12use std::time::Duration;
13use thiserror::Error;
14use tokio::time::{sleep, timeout};
15
16#[derive(Debug, Error)]
17pub enum TaskErrorReason {
18    #[error("failed to get task: {0}")]
19    Get(#[source] ServiceError),
20    #[error("failed to update task: {0}")]
21    Update(#[source] ServiceError),
22    #[error("expired before completion")]
23    Expired,
24    #[error("repeated input requests that were already answered")]
25    RepeatedInput,
26    #[error("failed: {error}")]
27    Failed { error: serde_json::Value },
28    #[error("was cancelled")]
29    Cancelled,
30    #[error("returned a malformed result: {0}")]
31    MalformedResult(#[source] serde_json::Error),
32    #[error("returned a task payload this client does not support (status {status:?})")]
33    UnsupportedPayload { status: TaskStatus },
34}
35
36pub(super) fn stream_task_events<'a>(
37    client: &'a McpClient,
38    created: CreateTaskResult,
39    progress: impl Stream<Item = ProgressNotificationParam> + 'a,
40) -> impl Stream<Item = ToolCallEvent> + 'a {
41    let polled = with_progress(progress, poll_task(client, created.task.clone())).map(|event| match event {
42        Either::Left(progress) => ToolCallEvent::Progress(progress),
43        Either::Right(event) => event,
44    });
45    futures::stream::once(std::future::ready(ToolCallEvent::TaskCreated(created))).chain(polled)
46}
47
48pub(super) async fn cancel_server_task(client: &McpClient, task_id: &str) {
49    let server_name = client.name();
50    match timeout(Duration::from_secs(1), client.peer.cancel_task(CancelTaskParams::new(task_id))).await {
51        Ok(Ok(())) => {}
52        Ok(Err(error)) => {
53            tracing::warn!(server = %server_name, %task_id, "Failed to cancel abandoned MCP task: {error}");
54        }
55        Err(_) => tracing::warn!(server = %server_name, %task_id, "Timed out cancelling abandoned MCP task"),
56    }
57}
58
59fn poll_task(client: &McpClient, mut task: Task) -> impl Stream<Item = ToolCallEvent> + '_ {
60    stream! {
61        let mut answered = HashSet::new();
62
63        loop {
64            if is_task_expired(&task) {
65                yield failed(&task.task_id, TaskErrorReason::Expired);
66                return;
67            }
68
69            let detailed_task = match client.peer.get_task(GetTaskParams::new(task.task_id.clone())).await {
70                Ok(result) => result.task,
71                Err(error) => {
72                    yield abandoned(client, &task.task_id, TaskErrorReason::Get(error)).await;
73                    return;
74                }
75            };
76
77            task = detailed_task.task;
78            if !task.status.is_terminal() {
79                yield ToolCallEvent::TaskStatus(task.clone());
80            }
81
82            match detailed_task.payload {
83                TaskPayload::Working => {}
84                TaskPayload::InputRequired { input_requests } => {
85                    if let Err(error) = answer_inputs(client, &task.task_id, input_requests, &mut answered).await {
86                        cancel_server_task(client, &task.task_id).await;
87                        yield error.into();
88                        return;
89                    }
90                }
91                TaskPayload::Completed { result } => {
92                    let result = serde_json::from_value(serde_json::Value::Object(result))
93                        .map_err(|source| task_error(&task.task_id, TaskErrorReason::MalformedResult(source)));
94                    yield ToolCallEvent::Done(result);
95                    return;
96                }
97                TaskPayload::Failed { error } => {
98                    yield failed(&task.task_id, TaskErrorReason::Failed { error: serde_json::Value::Object(error) });
99                    return;
100                }
101                TaskPayload::Cancelled => {
102                    yield failed(&task.task_id, TaskErrorReason::Cancelled);
103                    return;
104                }
105                _ => {
106                    let reason = TaskErrorReason::UnsupportedPayload { status: task.status };
107                    yield abandoned(client, &task.task_id, reason).await;
108                    return;
109                }
110            }
111
112            sleep(task.poll_interval_ms.map_or(DEFAULT_POLL_INTERVAL, Duration::from_millis)).await;
113        }
114    }
115}
116
117const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(1);
118
119async fn answer_inputs(
120    client: &McpClient,
121    task_id: &str,
122    input_requests: InputRequests,
123    answered: &mut HashSet<String>,
124) -> Result<(), ToolCallError> {
125    if input_requests.keys().any(|key| answered.contains(key)) {
126        return Err(task_error(task_id, TaskErrorReason::RepeatedInput));
127    }
128    let (responses, _) = client.handler.elicit_inputs(input_requests).await?;
129    answered.extend(responses.keys().cloned());
130    client
131        .peer
132        .update_task(UpdateTaskParams::new(task_id, responses))
133        .await
134        .map_err(|error| task_error(task_id, TaskErrorReason::Update(error)))
135}
136
137async fn abandoned(client: &McpClient, task_id: &str, reason: TaskErrorReason) -> ToolCallEvent {
138    cancel_server_task(client, task_id).await;
139    failed(task_id, reason)
140}
141
142fn failed(task_id: &str, reason: TaskErrorReason) -> ToolCallEvent {
143    task_error(task_id, reason).into()
144}
145
146fn task_error(task_id: &str, reason: TaskErrorReason) -> ToolCallError {
147    ToolCallError::Task { task_id: task_id.to_string(), reason: Box::new(reason) }
148}
149
150fn is_task_expired(task: &Task) -> bool {
151    if task.status.is_terminal() {
152        return false;
153    }
154    let Some(ttl_ms) = task.ttl_ms else {
155        return false;
156    };
157    let Ok(created_at) = chrono::DateTime::parse_from_rfc3339(&task.created_at) else {
158        tracing::warn!(task_id = %task.task_id, created_at = %task.created_at, "Ignoring malformed MCP task creation timestamp");
159        return false;
160    };
161    let Ok(ttl_ms) = i64::try_from(ttl_ms) else {
162        return false;
163    };
164    created_at
165        .with_timezone(&chrono::Utc)
166        .checked_add_signed(chrono::Duration::milliseconds(ttl_ms))
167        .is_some_and(|expires_at| chrono::Utc::now() > expires_at)
168}
mcp_utils/client/task.rs

mcp_utils/client/
task.rs
