1use std::collections::HashMap;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::Instant;
5
6use parking_lot::{Mutex, RwLock};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
10use tokio::sync::{broadcast, mpsc, oneshot};
11use tokio::task::JoinHandle;
12use tokio_util::sync::CancellationToken;
13use tracing::{Instrument, debug, error, warn};
14
15use crate::{Error, ErrorKind, ProtocolErrorKind};
16
17pub(crate) type InlineResponseCallback =
27 Box<dyn FnOnce(&JsonRpcResponse) -> Result<(), Error> + Send + Sync>;
28
29struct PendingRequest {
32 sender: oneshot::Sender<JsonRpcResponse>,
33 inline_callback: Option<InlineResponseCallback>,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
38#[serde(rename_all = "camelCase")]
39pub struct JsonRpcRequest {
40 pub jsonrpc: String,
42 pub id: u64,
44 pub method: String,
46 #[serde(skip_serializing_if = "Option::is_none")]
48 pub params: Option<Value>,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct JsonRpcResponse {
55 pub jsonrpc: String,
57 pub id: u64,
59 #[serde(skip_serializing_if = "Option::is_none")]
61 pub result: Option<Value>,
62 #[serde(skip_serializing_if = "Option::is_none")]
64 pub error: Option<JsonRpcError>,
65}
66
67#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct JsonRpcError {
70 pub code: i32,
72 pub message: String,
74 #[serde(skip_serializing_if = "Option::is_none")]
76 pub data: Option<Value>,
77}
78
79pub mod error_codes {
81 pub const METHOD_NOT_FOUND: i32 = -32601;
83 pub const INVALID_PARAMS: i32 = -32602;
85 #[allow(dead_code, reason = "standard JSON-RPC code, reserved for future use")]
87 pub const INTERNAL_ERROR: i32 = -32603;
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
92#[serde(rename_all = "camelCase")]
93pub struct JsonRpcNotification {
94 pub jsonrpc: String,
96 pub method: String,
98 #[serde(skip_serializing_if = "Option::is_none")]
100 pub params: Option<Value>,
101}
102
103#[derive(Debug, Clone, Serialize)]
105pub enum JsonRpcMessage {
106 Request(JsonRpcRequest),
108 Response(JsonRpcResponse),
110 Notification(JsonRpcNotification),
112}
113
114impl<'de> Deserialize<'de> for JsonRpcMessage {
123 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
124 where
125 D: serde::Deserializer<'de>,
126 {
127 let mut value = Value::deserialize(deserializer)?;
128 let obj = value
129 .as_object_mut()
130 .ok_or_else(|| serde::de::Error::custom("expected a JSON object"))?;
131
132 let has_id = obj.contains_key("id");
133 let has_method = obj.contains_key("method");
134
135 let payload_key = if has_id && !has_method {
138 "result"
139 } else {
140 "params"
141 };
142 let payload = obj.remove(payload_key).filter(|value| !value.is_null());
143
144 if has_id && has_method {
145 JsonRpcRequest::deserialize(value)
146 .map(|mut request| {
147 request.params = payload;
148 JsonRpcMessage::Request(request)
149 })
150 .map_err(serde::de::Error::custom)
151 } else if has_id {
152 JsonRpcResponse::deserialize(value)
153 .map(|mut response| {
154 response.result = payload;
155 JsonRpcMessage::Response(response)
156 })
157 .map_err(serde::de::Error::custom)
158 } else {
159 JsonRpcNotification::deserialize(value)
160 .map(|mut notification| {
161 notification.params = payload;
162 JsonRpcMessage::Notification(notification)
163 })
164 .map_err(serde::de::Error::custom)
165 }
166 }
167}
168
169impl JsonRpcRequest {
170 pub fn new(id: u64, method: &str, params: Option<Value>) -> Self {
172 Self {
173 jsonrpc: "2.0".to_string(),
174 id,
175 method: method.to_string(),
176 params,
177 }
178 }
179}
180
181impl JsonRpcResponse {
182 #[allow(dead_code)]
184 pub fn is_error(&self) -> bool {
185 self.error.is_some()
186 }
187}
188
189const CONTENT_LENGTH_HEADER: &str = "Content-Length: ";
190
191fn repair_lone_surrogates(body: &[u8]) -> Option<Vec<u8>> {
196 fn hex_escape_at(body: &[u8], index: usize) -> Option<u16> {
197 let digits = body.get(index + 2..index + 6)?;
198 let text = std::str::from_utf8(digits).ok()?;
199 u16::from_str_radix(text, 16).ok()
200 }
201
202 let mut repaired = None;
203 let mut in_string = false;
204 let mut index = 0;
205
206 while index < body.len() {
207 let byte = body[index];
208
209 if !in_string {
210 in_string = byte == b'"';
211 index += 1;
212 continue;
213 }
214
215 match byte {
216 b'"' => {
217 in_string = false;
218 index += 1;
219 }
220 b'\\' if body.get(index + 1) != Some(&b'u') => index += 2,
223 b'\\' => {
224 let Some(unit) = hex_escape_at(body, index) else {
225 index += 2;
226 continue;
227 };
228
229 let is_pair = (0xD800..0xDC00).contains(&unit)
230 && body.get(index + 6) == Some(&b'\\')
231 && body.get(index + 7) == Some(&b'u')
232 && hex_escape_at(body, index + 6)
233 .is_some_and(|low| (0xDC00..0xE000).contains(&low));
234
235 if is_pair {
236 index += 12;
237 continue;
238 }
239
240 if (0xD800..0xE000).contains(&unit) {
241 let output = repaired.get_or_insert_with(|| body.to_vec());
242 output[index..index + 6].copy_from_slice(br"\ufffd");
243 }
244 index += 6;
245 }
246 _ => index += 1,
247 }
248 }
249
250 repaired
251}
252
253struct WriteCommand {
262 frame: Vec<u8>,
263 ack: oneshot::Sender<Result<(), std::io::Error>>,
264}
265
266pub struct JsonRpcClient {
277 request_id: AtomicU64,
278 write_tx: mpsc::UnboundedSender<WriteCommand>,
285 pending_requests: Arc<RwLock<HashMap<u64, PendingRequest>>>,
286 notification_tx: broadcast::Sender<JsonRpcNotification>,
287 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
288 connection_closed: CancellationToken,
289 read_task: Mutex<Option<JoinHandle<()>>>,
290 write_task: Mutex<Option<JoinHandle<()>>>,
291}
292
293impl JsonRpcClient {
294 pub fn new(
301 writer: impl AsyncWrite + Unpin + Send + 'static,
302 reader: impl AsyncRead + Unpin + Send + 'static,
303 notification_tx: broadcast::Sender<JsonRpcNotification>,
304 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
305 ) -> Self {
306 let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteCommand>();
307
308 let writer_span = tracing::error_span!("jsonrpc_write_loop");
309 let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span));
310
311 let client = Self {
312 request_id: AtomicU64::new(1),
313 write_tx,
314 pending_requests: Arc::new(RwLock::new(HashMap::new())),
315 notification_tx,
316 request_tx,
317 connection_closed: CancellationToken::new(),
318 read_task: Mutex::new(None),
319 write_task: Mutex::new(Some(write_task)),
320 };
321
322 let pending_requests = client.pending_requests.clone();
323 let notification_tx_clone = client.notification_tx.clone();
324 let request_tx_clone = client.request_tx.clone();
325 let connection_closed = client.connection_closed.clone();
326 let reader_span = tracing::error_span!("jsonrpc_read_loop");
327
328 let read_task = tokio::spawn(
329 async move {
330 Self::read_loop(
331 reader,
332 pending_requests,
333 notification_tx_clone,
334 request_tx_clone,
335 )
336 .await;
337 connection_closed.cancel();
338 }
339 .instrument(reader_span),
340 );
341 *client.read_task.lock() = Some(read_task);
342
343 client
344 }
345
346 pub(crate) fn force_close(&self) {
347 self.connection_closed.cancel();
348 if let Some(task) = self.read_task.lock().take() {
349 task.abort();
350 }
351 if let Some(task) = self.write_task.lock().take() {
352 task.abort();
353 }
354 self.pending_requests.write().clear();
355 }
356
357 pub(crate) fn connection_closed_token(&self) -> CancellationToken {
358 self.connection_closed.child_token()
359 }
360
361 async fn write_loop(
374 mut writer: impl AsyncWrite + Unpin + Send + 'static,
375 mut rx: mpsc::UnboundedReceiver<WriteCommand>,
376 ) {
377 while let Some(WriteCommand { frame, ack }) = rx.recv().await {
378 let result = async {
379 writer.write_all(&frame).await?;
380 writer.flush().await?;
381 Ok::<_, std::io::Error>(())
382 }
383 .await;
384
385 let _ = ack.send(result);
389 }
390 }
391
392 async fn read_loop(
393 reader: impl AsyncRead + Unpin + Send,
394 pending_requests: Arc<RwLock<HashMap<u64, PendingRequest>>>,
395 notification_tx: broadcast::Sender<JsonRpcNotification>,
396 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
397 ) {
398 let mut reader = BufReader::new(reader);
399
400 loop {
401 match Self::read_message(&mut reader).await {
402 Ok(Some(message)) => match message {
403 JsonRpcMessage::Response(mut response) => {
404 let id = response.id;
405 let pending = pending_requests.write().remove(&id);
406 if let Some(PendingRequest {
407 sender,
408 inline_callback,
409 }) = pending
410 {
411 if let Some(cb) = inline_callback
417 && response.error.is_none()
418 {
419 let cb_outcome =
420 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
421 cb(&response)
422 }));
423 match cb_outcome {
424 Ok(Ok(())) => {}
425 Ok(Err(error)) => {
426 response.result = None;
427 response.error = Some(JsonRpcError {
428 code: -32603,
429 message: error.to_string(),
430 data: None,
431 });
432 }
433 Err(panic) => {
434 let message = panic
435 .downcast_ref::<&'static str>()
436 .map(|s| (*s).to_string())
437 .or_else(|| panic.downcast_ref::<String>().cloned())
438 .unwrap_or_else(|| {
439 "inline response callback panicked".to_string()
440 });
441 response.result = None;
442 response.error = Some(JsonRpcError {
443 code: -32603,
444 message,
445 data: None,
446 });
447 }
448 }
449 }
450 if sender.send(response).is_err() {
451 warn!(request_id = %id, "failed to send response for request");
452 }
453 } else {
454 warn!(request_id = %id, "received response for unknown request id");
455 }
456 }
457 JsonRpcMessage::Notification(notification) => {
458 let _ = notification_tx.send(notification);
459 }
460 JsonRpcMessage::Request(request) => {
461 if request_tx.send(request).is_err() {
462 warn!("failed to forward JSON-RPC request, channel closed");
463 }
464 }
465 },
466 Ok(None) => {
467 break;
468 }
469 Err(e) => {
470 error!(error = %e, "error reading from CLI");
471 break;
472 }
473 }
474 }
475
476 let mut pending = pending_requests.write();
479 if !pending.is_empty() {
480 warn!(
481 count = pending.len(),
482 "draining pending requests after read loop exit"
483 );
484 pending.clear();
485 }
486 }
487
488 async fn read_message(
489 reader: &mut BufReader<impl AsyncRead + Unpin>,
490 ) -> Result<Option<JsonRpcMessage>, Error> {
491 let mut line = String::new();
492 let mut content_length = None;
493
494 loop {
495 line.clear();
496 if reader.read_line(&mut line).await? == 0 {
497 return Ok(None);
498 }
499
500 let trimmed = line.trim();
501 if trimmed.is_empty() {
502 break;
503 }
504
505 if let Some(value) = trimmed.strip_prefix(CONTENT_LENGTH_HEADER) {
506 content_length = Some(value.trim().parse::<usize>().map_err(|_| {
507 Error::from(ErrorKind::Protocol(
508 ProtocolErrorKind::InvalidContentLength(value.trim().to_string()),
509 ))
510 })?);
511 }
512 }
513
514 let Some(length) = content_length else {
515 return Err(ErrorKind::Protocol(ProtocolErrorKind::MissingContentLength).into());
516 };
517
518 let mut body = vec![0u8; length];
519 reader.read_exact(&mut body).await?;
520
521 match serde_json::from_slice::<JsonRpcMessage>(&body) {
522 Ok(message) => Ok(Some(message)),
523 Err(error) => {
524 match repair_lone_surrogates(&body)
527 .and_then(|repaired| serde_json::from_slice::<JsonRpcMessage>(&repaired).ok())
528 {
529 Some(message) => {
530 warn!(
531 error = %error,
532 length,
533 "recovered JSON-RPC frame containing unpaired UTF-16 surrogates"
534 );
535 Ok(Some(message))
536 }
537 None => Err(error.into()),
538 }
539 }
540 }
541 }
542
543 #[allow(dead_code, reason = "public API exported via crate::JsonRpcClient")]
554 pub async fn send_request(
555 &self,
556 method: &str,
557 params: Option<serde_json::Value>,
558 ) -> Result<JsonRpcResponse, Error> {
559 self.send_request_with_inline_callback(method, params, None)
560 .await
561 }
562
563 pub(crate) async fn send_request_with_inline_callback(
580 &self,
581 method: &str,
582 params: Option<serde_json::Value>,
583 inline_callback: Option<InlineResponseCallback>,
584 ) -> Result<JsonRpcResponse, Error> {
585 let request_start = Instant::now();
586 let id = self.request_id.fetch_add(1, Ordering::SeqCst);
587 let request = JsonRpcRequest::new(id, method, params);
588
589 let (tx, rx) = oneshot::channel();
590 self.pending_requests.write().insert(
591 id,
592 PendingRequest {
593 sender: tx,
594 inline_callback,
595 },
596 );
597
598 let mut guard = PendingGuard {
603 map: &self.pending_requests,
604 id,
605 armed: true,
606 };
607
608 if let Err(error) = self.write(&request).await {
612 warn!(
613 elapsed_ms = request_start.elapsed().as_millis(),
614 method = %method,
615 request_id = id,
616 status = "failed",
617 error = %error,
618 "JsonRpcClient::send_request JSON-RPC request finished"
619 );
620 return Err(error);
621 }
622
623 let response = match rx.await {
624 Ok(response) => response,
625 Err(_) => {
626 let error = ErrorKind::Protocol(ProtocolErrorKind::RequestCancelled).into();
627 warn!(
628 elapsed_ms = request_start.elapsed().as_millis(),
629 method = %method,
630 request_id = id,
631 status = "failed",
632 error = %error,
633 "JsonRpcClient::send_request JSON-RPC request finished"
634 );
635 return Err(error);
636 }
637 };
638 guard.disarm();
639 if let Some(error) = &response.error {
640 warn!(
641 elapsed_ms = request_start.elapsed().as_millis(),
642 method = %method,
643 request_id = id,
644 status = "failed",
645 code = error.code,
646 error = %error.message,
647 "JsonRpcClient::send_request JSON-RPC request finished"
648 );
649 } else {
650 debug!(
651 elapsed_ms = request_start.elapsed().as_millis(),
652 method = %method,
653 request_id = id,
654 status = "succeeded",
655 "JsonRpcClient::send_request JSON-RPC request finished"
656 );
657 }
658 Ok(response)
659 }
660
661 pub async fn write<T: serde::Serialize>(&self, message: &T) -> Result<(), Error> {
670 let body = serde_json::to_vec(message)?;
671 let mut frame = Vec::with_capacity(CONTENT_LENGTH_HEADER.len() + 16 + body.len() + 4);
672 frame.extend_from_slice(CONTENT_LENGTH_HEADER.as_bytes());
673 frame.extend_from_slice(body.len().to_string().as_bytes());
674 frame.extend_from_slice(b"\r\n\r\n");
675 frame.extend_from_slice(&body);
676
677 let (ack_tx, ack_rx) = oneshot::channel();
678 self.write_tx
679 .send(WriteCommand { frame, ack: ack_tx })
680 .map_err(|_| {
681 Error::from(std::io::Error::new(
682 std::io::ErrorKind::BrokenPipe,
683 "writer actor has shut down",
684 ))
685 })?;
686
687 match ack_rx.await {
688 Ok(Ok(())) => Ok(()),
689 Ok(Err(e)) => Err(Error::from(e)),
690 Err(_) => Err(Error::from(std::io::Error::new(
691 std::io::ErrorKind::BrokenPipe,
692 "writer actor dropped ack without responding",
693 ))),
694 }
695 }
696}
697
698struct PendingGuard<'a> {
702 map: &'a RwLock<HashMap<u64, PendingRequest>>,
703 id: u64,
704 armed: bool,
705}
706
707impl PendingGuard<'_> {
708 fn disarm(&mut self) {
709 self.armed = false;
710 }
711}
712
713impl Drop for PendingGuard<'_> {
714 fn drop(&mut self) {
715 if self.armed {
716 self.map.write().remove(&self.id);
717 }
718 }
719}
720
721#[cfg(test)]
722mod tests {
723 use super::*;
724
725 #[test]
726 fn deserialize_notification() {
727 let json = r#"{"jsonrpc":"2.0","method":"session.event","params":{"id":"e1"}}"#;
728 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
729 assert!(matches!(msg, JsonRpcMessage::Notification(n) if n.method == "session.event"));
730 }
731
732 #[test]
733 fn deserialize_request() {
734 let json =
735 r#"{"jsonrpc":"2.0","id":5,"method":"permission.request","params":{"kind":"shell"}}"#;
736 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
737 assert!(
738 matches!(msg, JsonRpcMessage::Request(r) if r.id == 5 && r.method == "permission.request")
739 );
740 }
741
742 #[test]
743 fn deserialize_response_with_result() {
744 let json = r#"{"jsonrpc":"2.0","id":3,"result":{"ok":true}}"#;
745 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
746 assert!(matches!(msg, JsonRpcMessage::Response(r) if r.id == 3 && !r.is_error()));
747 }
748
749 #[test]
750 fn deserialize_error_response() {
751 let json = r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32600,"message":"Invalid Request","data":{"nested":[1,{"reason":"invalid"}]}}}"#;
752 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
753 match msg {
754 JsonRpcMessage::Response(r) => {
755 assert!(r.is_error());
756 let err = r.error.unwrap();
757 assert_eq!(err.code, -32600);
758 assert_eq!(err.message, "Invalid Request");
759 assert_eq!(
760 err.data,
761 Some(serde_json::json!({"nested": [1, {"reason": "invalid"}]}))
762 );
763 }
764 other => panic!("expected Response, got {other:?}"),
765 }
766 }
767
768 #[test]
769 fn deserialize_rejects_non_object() {
770 let result = serde_json::from_str::<JsonRpcMessage>(r#""not an object""#);
771 assert!(result.is_err());
772 }
773
774 #[test]
775 fn deserialize_preserves_optional_payloads() {
776 for payload in [
777 None,
778 Some(Value::Null),
779 Some(serde_json::json!(false)),
780 Some(serde_json::json!(42)),
781 Some(serde_json::json!("text")),
782 Some(serde_json::json!([{"nested": [1, null, true]}])),
783 Some(serde_json::json!({"rows": [{"content": "result"}]})),
784 ] {
785 for mut envelope in [
786 serde_json::json!({"jsonrpc": "2.0", "method": "notify"}),
787 serde_json::json!({"jsonrpc": "2.0", "id": 1, "method": "request"}),
788 serde_json::json!({"jsonrpc": "2.0", "id": 1}),
789 ] {
790 let (payload_key, ignored_key) = if envelope.get("method").is_some() {
791 ("params", "result")
792 } else {
793 ("result", "params")
794 };
795 envelope[ignored_key] = serde_json::json!({"ignored": "opposite payload"});
796 if let Some(payload) = &payload {
797 envelope[payload_key] = payload.clone();
798 }
799 let actual = match serde_json::from_value::<JsonRpcMessage>(envelope).unwrap() {
800 JsonRpcMessage::Request(request) => request.params,
801 JsonRpcMessage::Response(response) => response.result,
802 JsonRpcMessage::Notification(notification) => notification.params,
803 };
804 assert_eq!(actual, payload.clone().filter(|value| !value.is_null()));
805 }
806 }
807 }
808
809 #[test]
810 fn deserialize_rejects_invalid_metadata() {
811 for json in [
812 r#"{"jsonrpc":null,"method":"notify","params":{"nested":[1]}}"#,
813 r#"{"jsonrpc":"2.0","method":42,"params":{"nested":[1]}}"#,
814 r#"{"jsonrpc":"2.0","id":null,"result":{}}"#,
815 r#"{"jsonrpc":"2.0","id":"1","result":{}}"#,
816 r#"{"jsonrpc":"2.0","id":-1,"result":{}}"#,
817 r#"{"jsonrpc":"2.0","id":1,"method":null,"params":{}}"#,
818 r#"{"jsonrpc":"2.0","id":1,"result":{},"error":{"code":"bad","message":"error"}}"#,
819 ] {
820 assert!(
821 serde_json::from_str::<JsonRpcMessage>(json).is_err(),
822 "{json}"
823 );
824 }
825 }
826
827 #[test]
828 fn request_new_sets_version() {
829 let req = JsonRpcRequest::new(42, "test.method", None);
830 assert_eq!(req.jsonrpc, "2.0");
831 assert_eq!(req.id, 42);
832 assert_eq!(req.method, "test.method");
833 assert!(req.params.is_none());
834 }
835
836 #[test]
837 fn request_serializes_camel_case() {
838 let req = JsonRpcRequest::new(1, "ping", Some(serde_json::json!({})));
839 let json = serde_json::to_string(&req).unwrap();
840 assert!(json.contains(r#""jsonrpc":"2.0""#));
841 assert!(json.contains(r#""id":1"#));
842 assert!(json.contains(r#""method":"ping""#));
843 }
844
845 #[test]
846 fn notification_without_params_omits_field() {
847 let n = JsonRpcNotification {
848 jsonrpc: "2.0".into(),
849 method: "ping".into(),
850 params: None,
851 };
852 let json = serde_json::to_string(&n).unwrap();
853 assert!(!json.contains("params"));
854 }
855
856 #[test]
857 fn response_without_error_omits_field() {
858 let r = JsonRpcResponse {
859 jsonrpc: "2.0".into(),
860 id: 1,
861 result: Some(serde_json::json!(true)),
862 error: None,
863 };
864 let json = serde_json::to_string(&r).unwrap();
865 assert!(!json.contains("error"));
866 }
867}