Skip to main content

mcp_streamable_proxy/proxy_handler/
raw_dispatch.rs

1use serde::Serialize;
2use serde::de::DeserializeOwned;
3use serde_json::Value;
4
5use super::*;
6
7type JsonResult = Result<Value, String>;
8
9impl ProxyHandler {
10    /// Call a backend method through the version-independent JSON bridge.
11    pub async fn call_peer_method(&self, method: &str, params: Value) -> JsonResult {
12        let peer = self.raw_peer()?;
13        match method {
14            "tools/list" => list_tools(peer, params).await,
15            "tools/call" => call_tool(peer, params).await,
16            "resources/list" => list_resources(peer, params).await,
17            "resources/read" => read_resource(peer, params).await,
18            "resources/templates/list" => list_resource_templates(peer, params).await,
19            "resources/subscribe" => subscribe(peer, params).await,
20            "resources/unsubscribe" => unsubscribe(peer, params).await,
21            "prompts/list" => list_prompts(peer, params).await,
22            "prompts/get" => get_prompt(peer, params).await,
23            "completion/complete" => complete(peer, params).await,
24            "logging/setLevel" => set_log_level(peer, params).await,
25            "tasks/list" => list_tasks(peer, params).await,
26            "tasks/get" => get_task(peer, params).await,
27            "tasks/result" => get_task_result(peer, params).await,
28            "tasks/cancel" => cancel_task(peer, params).await,
29            _ => Err(format!("unsupported method: {method}")),
30        }
31    }
32
33    fn raw_peer(&self) -> Result<Peer<RoleClient>, String> {
34        let inner = self
35            .peer
36            .load_full()
37            .ok_or_else(|| "backend connection is not available (reconnecting)".to_string())?;
38        if inner.peer.is_transport_closed() {
39            return Err("backend transport is closed".to_string());
40        }
41        Ok(inner.peer.clone())
42    }
43}
44
45fn parse<T: DeserializeOwned>(params: Value, method: &str) -> Result<T, String> {
46    serde_json::from_value(params).map_err(|error| format!("invalid params for {method}: {error}"))
47}
48
49fn serialize<T: Serialize>(value: T) -> JsonResult {
50    serde_json::to_value(value)
51        .map_err(|error| format!("failed to serialize backend result: {error}"))
52}
53
54fn page_params(
55    params: Value,
56    method: &str,
57) -> Result<Option<rmcp::model::PaginatedRequestParams>, String> {
58    if params.is_null() {
59        return Ok(None);
60    }
61    parse(params, method)
62}
63
64async fn list_tools(peer: Peer<RoleClient>, params: Value) -> JsonResult {
65    let result = peer
66        .list_tools(page_params(params, "tools/list")?)
67        .await
68        .map_err(|error| format!("tools/list failed: {error:?}"))?;
69    serialize(result)
70}
71
72async fn call_tool(peer: Peer<RoleClient>, params: Value) -> JsonResult {
73    let request: CallToolRequestParams = parse(params, "tools/call")?;
74    if request.task.is_some() {
75        let result = peer
76            .send_request(ClientRequest::CallToolRequest(CallToolRequest::new(
77                request,
78            )))
79            .await
80            .map_err(|error| format!("tools/call task enqueue failed: {error:?}"))?;
81        return match result {
82            ServerResult::CreateTaskResult(result) => serialize(result),
83            other => Err(format!("unexpected tools/call task response: {other:?}")),
84        };
85    }
86
87    match peer
88        .call_tool_once(request)
89        .await
90        .map_err(|error| format!("tools/call failed: {error:?}"))?
91    {
92        CallToolResponse::Complete(result) => serialize(result),
93        CallToolResponse::InputRequired(result) => serialize(result),
94        other => Err(format!("unsupported tools/call response: {other:?}")),
95    }
96}
97
98async fn list_resources(peer: Peer<RoleClient>, params: Value) -> JsonResult {
99    let result = peer
100        .list_resources(page_params(params, "resources/list")?)
101        .await
102        .map_err(|error| format!("resources/list failed: {error:?}"))?;
103    serialize(result)
104}
105
106async fn read_resource(peer: Peer<RoleClient>, params: Value) -> JsonResult {
107    let request = parse(params, "resources/read")?;
108    match peer
109        .read_resource_once(request)
110        .await
111        .map_err(|error| format!("resources/read failed: {error:?}"))?
112    {
113        ReadResourceResponse::Complete(result) => serialize(result),
114        ReadResourceResponse::InputRequired(result) => serialize(result),
115        other => Err(format!("unsupported resources/read response: {other:?}")),
116    }
117}
118
119async fn list_resource_templates(peer: Peer<RoleClient>, params: Value) -> JsonResult {
120    let result = peer
121        .list_resource_templates(page_params(params, "resources/templates/list")?)
122        .await
123        .map_err(|error| format!("resources/templates/list failed: {error:?}"))?;
124    serialize(result)
125}
126
127async fn subscribe(peer: Peer<RoleClient>, params: Value) -> JsonResult {
128    peer.subscribe(parse(params, "resources/subscribe")?)
129        .await
130        .map_err(|error| format!("resources/subscribe failed: {error:?}"))?;
131    Ok(serde_json::json!({}))
132}
133
134async fn unsubscribe(peer: Peer<RoleClient>, params: Value) -> JsonResult {
135    peer.unsubscribe(parse(params, "resources/unsubscribe")?)
136        .await
137        .map_err(|error| format!("resources/unsubscribe failed: {error:?}"))?;
138    Ok(serde_json::json!({}))
139}
140
141async fn list_prompts(peer: Peer<RoleClient>, params: Value) -> JsonResult {
142    let result = peer
143        .list_prompts(page_params(params, "prompts/list")?)
144        .await
145        .map_err(|error| format!("prompts/list failed: {error:?}"))?;
146    serialize(result)
147}
148
149async fn get_prompt(peer: Peer<RoleClient>, params: Value) -> JsonResult {
150    let request = parse(params, "prompts/get")?;
151    match peer
152        .get_prompt_once(request)
153        .await
154        .map_err(|error| format!("prompts/get failed: {error:?}"))?
155    {
156        GetPromptResponse::Complete(result) => serialize(result),
157        GetPromptResponse::InputRequired(result) => serialize(result),
158        other => Err(format!("unsupported prompts/get response: {other:?}")),
159    }
160}
161
162async fn complete(peer: Peer<RoleClient>, params: Value) -> JsonResult {
163    let result = peer
164        .complete(parse(params, "completion/complete")?)
165        .await
166        .map_err(|error| format!("completion/complete failed: {error:?}"))?;
167    serialize(result)
168}
169
170async fn set_log_level(peer: Peer<RoleClient>, params: Value) -> JsonResult {
171    #[allow(deprecated)]
172    peer.set_level(parse(params, "logging/setLevel")?)
173        .await
174        .map_err(|error| format!("logging/setLevel failed: {error:?}"))?;
175    Ok(serde_json::json!({}))
176}
177
178async fn list_tasks(peer: Peer<RoleClient>, params: Value) -> JsonResult {
179    let result = peer
180        .send_request(ClientRequest::ListTasksRequest(ListTasksRequest {
181            method: Default::default(),
182            params: page_params(params, "tasks/list")?,
183            extensions: Default::default(),
184        }))
185        .await
186        .map_err(|error| format!("tasks/list failed: {error:?}"))?;
187    match result {
188        ServerResult::ListTasksResult(result) => serialize(result),
189        other => Err(format!("unexpected tasks/list response: {other:?}")),
190    }
191}
192
193async fn get_task(peer: Peer<RoleClient>, params: Value) -> JsonResult {
194    let request = GetTaskRequest::new(parse(params, "tasks/get")?);
195    let result = peer
196        .send_request(ClientRequest::GetTaskRequest(request))
197        .await
198        .map_err(|error| format!("tasks/get failed: {error:?}"))?;
199    match result {
200        ServerResult::GetTaskResult(result) => serialize(result),
201        other => Err(format!("unexpected tasks/get response: {other:?}")),
202    }
203}
204
205async fn get_task_result(peer: Peer<RoleClient>, params: Value) -> JsonResult {
206    let request = GetTaskPayloadRequest::new(parse(params, "tasks/result")?);
207    let result = peer
208        .send_request(ClientRequest::GetTaskPayloadRequest(request))
209        .await
210        .map_err(|error| format!("tasks/result failed: {error:?}"))?;
211    match result {
212        ServerResult::GetTaskPayloadResult(result) => serialize(result),
213        other => Err(format!("unexpected tasks/result response: {other:?}")),
214    }
215}
216
217async fn cancel_task(peer: Peer<RoleClient>, params: Value) -> JsonResult {
218    let request = CancelTaskRequest::new(parse(params, "tasks/cancel")?);
219    let result = peer
220        .send_request(ClientRequest::CancelTaskRequest(request))
221        .await
222        .map_err(|error| format!("tasks/cancel failed: {error:?}"))?;
223    match result {
224        ServerResult::CancelTaskResult(result) => serialize(result),
225        other => Err(format!("unexpected tasks/cancel response: {other:?}")),
226    }
227}