1use std::collections::HashMap;
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12use serde::Serialize;
13
14use crate::error::Error;
15use crate::types::{Message, RequestId, Response};
16
17type BoxedHandler = Box<
19 dyn Fn(
20 serde_json::Value,
21 ) -> Pin<Box<dyn Future<Output = Result<serde_json::Value, Error>> + Send>>
22 + Send
23 + Sync,
24>;
25
26pub struct JsonRpc {
49 handlers: HashMap<String, BoxedHandler>,
50}
51
52impl JsonRpc {
53 pub fn new() -> Self {
55 Self {
56 handlers: HashMap::new(),
57 }
58 }
59
60 pub fn add<F, P, R, Fut>(mut self, method: &str, handler: F) -> Self
79 where
80 F: Fn(P) -> Fut + Send + Sync + 'static,
81 Fut: Future<Output = Result<R, Error>> + Send + Sync + 'static,
82 P: serde::de::DeserializeOwned + Send + Sync + 'static,
83 R: Serialize + Send + Sync + 'static,
84 {
85 let handler = Arc::new(handler);
86 let boxed: BoxedHandler = Box::new(move |params: serde_json::Value| {
87 let handler = Arc::clone(&handler);
88 Box::pin(async move {
89 let parsed: P = serde_json::from_value(params)?;
90 let result = handler(parsed).await?;
91 Ok(serde_json::to_value(result)?)
92 })
93 });
94
95 self.handlers.insert(method.to_string(), boxed);
96 self
97 }
98
99 pub async fn call(&self, json_str: &str) -> Option<String> {
111 let value: serde_json::Value = match serde_json::from_str(json_str) {
112 Ok(v) => v,
113 Err(_) => {
114 let error = crate::types::Error::parse_error("Parse error");
115 let response = Response::error(RequestId::Null, error);
116 match serde_json::to_string(&response) {
117 Ok(s) => return Some(s),
118 Err(e) => {
119 tracing::error!("Failed to serialize parse error response: {}", e);
120 return None;
121 }
122 }
123 }
124 };
125
126 let request_id = value.get("id").and_then(|id_value| match id_value {
127 serde_json::Value::Null => Some(RequestId::Null),
128 serde_json::Value::Number(n) => n.as_u64().map(RequestId::Number),
129 serde_json::Value::String(s) => Some(RequestId::String(s.clone())),
130 _ => None,
131 });
132
133 let message = match Message::from_json(value) {
134 Ok(msg) => msg,
135 Err(Error::InvalidRequest(_)) => {
136 let error = crate::types::Error::invalid_request("Invalid Request");
137 let id_to_use = request_id.unwrap_or(RequestId::Null);
138 let response = Response::error(id_to_use, error);
139 match serde_json::to_string(&response) {
140 Ok(s) => return Some(s),
141 Err(e) => {
142 tracing::error!("Failed to serialize invalid request response: {}", e);
143 return None;
144 }
145 }
146 }
147 Err(_) => {
148 let error = crate::types::Error::internal_error("Internal error");
149 let response = Response::error(request_id.unwrap_or(RequestId::Null), error);
150 match serde_json::to_string(&response) {
151 Ok(s) => return Some(s),
152 Err(e) => {
153 tracing::error!("Failed to serialize internal error response: {}", e);
154 return None;
155 }
156 }
157 }
158 };
159
160 match message {
161 Message::Request(request) => {
162 let method_name = &request.method;
163 let params = request.params.unwrap_or(serde_json::Value::Null);
164 let request_id = request.id.clone();
165 let response = if let Some(handler) = self.handlers.get(method_name) {
166 let result = handler(params).await;
167 match result {
168 Ok(result_value) => Response::success(request_id, result_value),
169 Err(e) => {
170 let error = match e {
171 crate::error::Error::RpcError { code, message } => {
172 crate::types::Error::new(code, message, None)
173 }
174 _ => crate::types::Error::new(-32603, e.to_string(), None),
175 };
176 Response::error(request_id, error)
177 }
178 }
179 } else {
180 let error = crate::types::Error::method_not_found(format!(
181 "Unknown method: {}",
182 method_name
183 ));
184 Response::error(request_id, error)
185 };
186 match serde_json::to_string(&response) {
187 Ok(s) => Some(s),
188 Err(e) => {
189 tracing::error!("Failed to serialize response: {}", e);
190 None
191 }
192 }
193 }
194 Message::Notification(notification) => {
195 if let Some(handler) = self.handlers.get(¬ification.method) {
196 let params = notification.params.unwrap_or(serde_json::Value::Null);
197 let _ = handler(params).await;
198 }
199 None
200 }
201 Message::Batch(messages) => {
202 let mut responses = Vec::new();
203
204 for message in messages {
205 match message {
206 Message::Request(request) => {
207 let method_name = &request.method;
208 let params = request.params.unwrap_or(serde_json::Value::Null);
209 let id = request.id;
210 let response = if let Some(handler) = self.handlers.get(method_name) {
211 let result = handler(params).await;
212 match result {
213 Ok(result_value) => Response::success(id, result_value),
214 Err(e) => {
215 let error = match e {
216 crate::error::Error::RpcError { code, message } => {
217 crate::types::Error::new(code, message, None)
218 }
219 _ => crate::types::Error::new(
220 -32603,
221 e.to_string(),
222 None,
223 ),
224 };
225 Response::error(id, error)
226 }
227 }
228 } else {
229 let error = crate::types::Error::method_not_found(format!(
230 "Unknown method: {}",
231 method_name
232 ));
233 Response::error(id, error)
234 };
235 responses.push(response);
236 }
237 Message::Notification(notification) => {
238 if let Some(handler) = self.handlers.get(¬ification.method) {
239 let params = notification.params.unwrap_or(serde_json::Value::Null);
240 let _ = handler(params).await;
241 }
242 }
243 Message::Response(response) => {
244 responses.push(response);
245 }
246 Message::Batch(_) => {
247 let error_response = Response::error(
248 crate::types::RequestId::Null,
249 crate::types::Error::invalid_request("Invalid Request"),
250 );
251 responses.push(error_response);
252 }
253 }
254 }
255
256 match serde_json::to_string(&responses) {
257 Ok(s) => Some(s),
258 Err(e) => {
259 tracing::error!("Failed to serialize batch responses: {}", e);
260 None
261 }
262 }
263 }
264 Message::Response(_response) => None,
265 }
266 }
267}