mcp_streamable_proxy/proxy_handler/
raw_dispatch.rs1use serde::Serialize;
2use serde::de::DeserializeOwned;
3use serde_json::Value;
4
5use super::*;
6
7type JsonResult = Result<Value, String>;
8
9impl ProxyHandler {
10 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 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
208async fn forward_raw(peer: Peer<RoleClient>, method: &str, params: Value) -> JsonResult {
215 let result = peer
216 .send_request(ClientRequest::CustomRequest(
217 rmcp::model::CustomRequest::new(method.to_string(), Some(params)),
218 ))
219 .await
220 .map_err(|error| format!("{method} failed: {error:?}"))?;
221 match result {
222 ServerResult::CustomResult(result) => Ok(result.0),
223 other => Err(format!("unexpected {method} response: {other:?}")),
224 }
225}