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" => forward_raw(peer, "tasks/list", params).await,
26            "tasks/get" => get_task(peer, params).await,
27            "tasks/result" => forward_raw(peer, "tasks/result", params).await,
28            "tasks/update" => update_task(peer, params).await,
29            "tasks/cancel" => cancel_task(peer, params).await,
30            _ => Err(format!("unsupported method: {method}")),
31        }
32    }
33
34    fn raw_peer(&self) -> Result<Peer<RoleClient>, String> {
35        let inner = self
36            .peer
37            .load_full()
38            .ok_or_else(|| "backend connection is not available (reconnecting)".to_string())?;
39        if inner.peer.is_transport_closed() {
40            return Err("backend transport is closed".to_string());
41        }
42        Ok(inner.peer.clone())
43    }
44
45    /// Forward an arbitrary JSON-RPC method to the backend verbatim (transparent bridge).
46    ///
47    /// Unlike [`call_peer_method`](Self::call_peer_method) (a whitelist dispatcher),
48    /// this forwards ANY method string to the backend as a custom request, so the
49    /// proxy stays transparent for legacy methods (`tasks/list`, `tasks/result`) and
50    /// vendor extensions alike. Used by `ServerHandler::on_custom_request`.
51    pub(super) async fn forward_to_backend(&self, method: &str, params: Value) -> JsonResult {
52        let peer = self.raw_peer()?;
53        forward_raw(peer, method, params).await
54    }
55}
56
57fn parse<T: DeserializeOwned>(params: Value, method: &str) -> Result<T, String> {
58    serde_json::from_value(params).map_err(|error| format!("invalid params for {method}: {error}"))
59}
60
61fn serialize<T: Serialize>(value: T) -> JsonResult {
62    serde_json::to_value(value)
63        .map_err(|error| format!("failed to serialize backend result: {error}"))
64}
65
66fn page_params(
67    params: Value,
68    method: &str,
69) -> Result<Option<rmcp::model::PaginatedRequestParams>, String> {
70    if params.is_null() {
71        return Ok(None);
72    }
73    parse(params, method)
74}
75
76async fn list_tools(peer: Peer<RoleClient>, params: Value) -> JsonResult {
77    let result = peer
78        .list_tools(page_params(params, "tools/list")?)
79        .await
80        .map_err(|error| format!("tools/list failed: {error:?}"))?;
81    serialize(result)
82}
83
84async fn call_tool(peer: Peer<RoleClient>, params: Value) -> JsonResult {
85    let request: CallToolRequestParams = parse(params, "tools/call")?;
86    match peer
87        .call_tool_once(request)
88        .await
89        .map_err(|error| format!("tools/call failed: {error:?}"))?
90    {
91        CallToolResponse::Complete(result) => serialize(result),
92        CallToolResponse::InputRequired(result) => serialize(result),
93        CallToolResponse::Task(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
127#[allow(deprecated)]
128async fn subscribe(peer: Peer<RoleClient>, params: Value) -> JsonResult {
129    peer.subscribe(parse(params, "resources/subscribe")?)
130        .await
131        .map_err(|error| format!("resources/subscribe failed: {error:?}"))?;
132    Ok(serde_json::json!({}))
133}
134
135#[allow(deprecated)]
136async fn unsubscribe(peer: Peer<RoleClient>, params: Value) -> JsonResult {
137    peer.unsubscribe(parse(params, "resources/unsubscribe")?)
138        .await
139        .map_err(|error| format!("resources/unsubscribe failed: {error:?}"))?;
140    Ok(serde_json::json!({}))
141}
142
143async fn list_prompts(peer: Peer<RoleClient>, params: Value) -> JsonResult {
144    let result = peer
145        .list_prompts(page_params(params, "prompts/list")?)
146        .await
147        .map_err(|error| format!("prompts/list failed: {error:?}"))?;
148    serialize(result)
149}
150
151async fn get_prompt(peer: Peer<RoleClient>, params: Value) -> JsonResult {
152    let request = parse(params, "prompts/get")?;
153    match peer
154        .get_prompt_once(request)
155        .await
156        .map_err(|error| format!("prompts/get failed: {error:?}"))?
157    {
158        GetPromptResponse::Complete(result) => serialize(result),
159        GetPromptResponse::InputRequired(result) => serialize(result),
160        other => Err(format!("unsupported prompts/get response: {other:?}")),
161    }
162}
163
164async fn complete(peer: Peer<RoleClient>, params: Value) -> JsonResult {
165    let result = peer
166        .complete(parse(params, "completion/complete")?)
167        .await
168        .map_err(|error| format!("completion/complete failed: {error:?}"))?;
169    serialize(result)
170}
171
172async fn set_log_level(peer: Peer<RoleClient>, params: Value) -> JsonResult {
173    #[allow(deprecated)]
174    peer.set_level(parse(params, "logging/setLevel")?)
175        .await
176        .map_err(|error| format!("logging/setLevel failed: {error:?}"))?;
177    Ok(serde_json::json!({}))
178}
179
180async fn get_task(peer: Peer<RoleClient>, params: Value) -> JsonResult {
181    let request = GetTaskRequest::new(parse(params, "tasks/get")?);
182    let result = peer
183        .send_request(ClientRequest::GetTaskRequest(request))
184        .await
185        .map_err(|error| format!("tasks/get failed: {error:?}"))?;
186    match result {
187        ServerResult::GetTaskResult(result) => serialize(result),
188        other => Err(format!("unexpected tasks/get response: {other:?}")),
189    }
190}
191
192async fn update_task(peer: Peer<RoleClient>, params: Value) -> JsonResult {
193    let request: rmcp::model::UpdateTaskParams = parse(params, "tasks/update")?;
194    peer.update_task(request)
195        .await
196        .map_err(|error| format!("tasks/update failed: {error:?}"))?;
197    Ok(serde_json::json!({}))
198}
199
200async fn cancel_task(peer: Peer<RoleClient>, params: Value) -> JsonResult {
201    let request: rmcp::model::CancelTaskParams = parse(params, "tasks/cancel")?;
202    peer.cancel_task(request)
203        .await
204        .map_err(|error| format!("tasks/cancel failed: {error:?}"))?;
205    Ok(serde_json::json!({}))
206}
207
208/// Forward a method to the backend verbatim as a custom JSON-RPC request.
209///
210/// Used for legacy methods (`tasks/list`, `tasks/result`) whose typed
211/// request/result were dropped in rmcp 3.1.0 (SEP-2663 Tasks redesign) but
212/// that older backends may still implement. Keeps the proxy a transparent,
213/// version-tolerant bridge instead of hard-failing on spec-removed methods.
214async fn forward_raw(peer: Peer<RoleClient>, method: &str, params: Value) -> JsonResult {
215    let result = peer
216        .send_request(ClientRequest::CustomRequest(rmcp::model::CustomRequest::new(
217            method.to_string(),
218            Some(params),
219        )))
220        .await
221        .map_err(|error| format!("{method} failed: {error:?}"))?;
222    match result {
223        ServerResult::CustomResult(result) => Ok(result.0),
224        other => Err(format!("unexpected {method} response: {other:?}")),
225    }
226}