soma_mcp_client/upstream/pool/
tasks.rs

1use std::collections::BTreeMap;
2
3use rmcp::model::{CancelTaskParams, GetTaskParams, UpdateTaskParams};
4use serde_json::Value;
5
6use crate::upstream::{CapScope, UpstreamError};
7
8use super::{UpstreamPool, ensure_routable};
9
10impl UpstreamPool {
11    pub async fn get_task(&self, upstream: &str, task_id: &str) -> Result<Value, UpstreamError> {
12        self.ensure_connected(upstream).await?;
13        let peer = self.task_peer(upstream)?;
14        let result = peer
15            .get_task(GetTaskParams::new(task_id))
16            .await
17            .map_err(|error| task_error(upstream, "tasks/get", error))?;
18        let value = serde_json::to_value(result)
19            .map_err(|error| task_error(upstream, "tasks/get", error))?;
20        let bytes = serde_json::to_vec(&value).map_or(usize::MAX, |bytes| bytes.len());
21        self.response_caps().enforce(CapScope::ToolsCall, bytes)?;
22        Ok(value)
23    }
24
25    pub async fn update_task(
26        &self,
27        upstream: &str,
28        task_id: &str,
29        input_responses: BTreeMap<String, Value>,
30    ) -> Result<(), UpstreamError> {
31        self.ensure_connected(upstream).await?;
32        let peer = self.task_peer(upstream)?;
33        peer.update_task(UpdateTaskParams::new(task_id, input_responses))
34            .await
35            .map_err(|error| task_error(upstream, "tasks/update", error))
36    }
37
38    pub async fn cancel_task(&self, upstream: &str, task_id: &str) -> Result<(), UpstreamError> {
39        self.ensure_connected(upstream).await?;
40        let peer = self.task_peer(upstream)?;
41        peer.cancel_task(CancelTaskParams::new(task_id))
42            .await
43            .map_err(|error| task_error(upstream, "tasks/cancel", error))
44    }
45
46    fn task_peer(
47        &self,
48        upstream: &str,
49    ) -> Result<rmcp::service::Peer<rmcp::RoleClient>, UpstreamError> {
50        self.with_entry(upstream, |entry| {
51            ensure_routable(entry)?;
52            entry
53                .live
54                .as_ref()
55                .map(|live| live.peer())
56                .ok_or_else(|| UpstreamError::Unsupported {
57                    upstream: upstream.to_owned(),
58                    capability: "tasks",
59                })
60        })
61    }
62}
63
64fn task_error(
65    upstream: &str,
66    operation: &'static str,
67    error: impl std::fmt::Display,
68) -> UpstreamError {
69    UpstreamError::LiveCall {
70        upstream: upstream.to_owned(),
71        operation,
72        message: error.to_string(),
73    }
74}
75
76#[cfg(test)]
77#[path = "tasks_tests.rs"]
78mod tests;
soma_mcp_client/upstream/pool/tasks.rs

soma_mcp_client/upstream/pool/
tasks.rs
