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 tracing::{Instrument, debug, error, warn};
13
14use crate::{Error, ErrorKind, ProtocolErrorKind};
15
16pub(crate) type InlineResponseCallback =
26 Box<dyn FnOnce(&JsonRpcResponse) -> Result<(), Error> + Send + Sync>;
27
28struct PendingRequest {
31 sender: oneshot::Sender<JsonRpcResponse>,
32 inline_callback: Option<InlineResponseCallback>,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
37#[serde(rename_all = "camelCase")]
38pub struct JsonRpcRequest {
39 pub jsonrpc: String,
41 pub id: u64,
43 pub method: String,
45 #[serde(skip_serializing_if = "Option::is_none")]
47 pub params: Option<Value>,
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize)]
52#[serde(rename_all = "camelCase")]
53pub struct JsonRpcResponse {
54 pub jsonrpc: String,
56 pub id: u64,
58 #[serde(skip_serializing_if = "Option::is_none")]
60 pub result: Option<Value>,
61 #[serde(skip_serializing_if = "Option::is_none")]
63 pub error: Option<JsonRpcError>,
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct JsonRpcError {
69 pub code: i32,
71 pub message: String,
73 #[serde(skip_serializing_if = "Option::is_none")]
75 pub data: Option<Value>,
76}
77
78pub mod error_codes {
80 pub const METHOD_NOT_FOUND: i32 = -32601;
82 pub const INVALID_PARAMS: i32 = -32602;
84 #[allow(dead_code, reason = "standard JSON-RPC code, reserved for future use")]
86 pub const INTERNAL_ERROR: i32 = -32603;
87}
88
89#[derive(Debug, Clone, Serialize, Deserialize)]
91#[serde(rename_all = "camelCase")]
92pub struct JsonRpcNotification {
93 pub jsonrpc: String,
95 pub method: String,
97 #[serde(skip_serializing_if = "Option::is_none")]
99 pub params: Option<Value>,
100}
101
102#[derive(Debug, Clone, Serialize)]
104pub enum JsonRpcMessage {
105 Request(JsonRpcRequest),
107 Response(JsonRpcResponse),
109 Notification(JsonRpcNotification),
111}
112
113impl<'de> Deserialize<'de> for JsonRpcMessage {
122 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
123 where
124 D: serde::Deserializer<'de>,
125 {
126 let value = Value::deserialize(deserializer)?;
127 let obj = value
128 .as_object()
129 .ok_or_else(|| serde::de::Error::custom("expected a JSON object"))?;
130
131 let has_id = obj.contains_key("id");
132 let has_method = obj.contains_key("method");
133
134 if has_id && has_method {
135 JsonRpcRequest::deserialize(value)
136 .map(JsonRpcMessage::Request)
137 .map_err(serde::de::Error::custom)
138 } else if has_id {
139 JsonRpcResponse::deserialize(value)
140 .map(JsonRpcMessage::Response)
141 .map_err(serde::de::Error::custom)
142 } else {
143 JsonRpcNotification::deserialize(value)
144 .map(JsonRpcMessage::Notification)
145 .map_err(serde::de::Error::custom)
146 }
147 }
148}
149
150impl JsonRpcRequest {
151 pub fn new(id: u64, method: &str, params: Option<Value>) -> Self {
153 Self {
154 jsonrpc: "2.0".to_string(),
155 id,
156 method: method.to_string(),
157 params,
158 }
159 }
160}
161
162impl JsonRpcResponse {
163 #[allow(dead_code)]
165 pub fn is_error(&self) -> bool {
166 self.error.is_some()
167 }
168}
169
170const CONTENT_LENGTH_HEADER: &str = "Content-Length: ";
171
172fn repair_lone_surrogates(body: &[u8]) -> Option<Vec<u8>> {
177 fn hex_escape_at(body: &[u8], index: usize) -> Option<u16> {
178 let digits = body.get(index + 2..index + 6)?;
179 let text = std::str::from_utf8(digits).ok()?;
180 u16::from_str_radix(text, 16).ok()
181 }
182
183 let mut repaired = None;
184 let mut in_string = false;
185 let mut index = 0;
186
187 while index < body.len() {
188 let byte = body[index];
189
190 if !in_string {
191 in_string = byte == b'"';
192 index += 1;
193 continue;
194 }
195
196 match byte {
197 b'"' => {
198 in_string = false;
199 index += 1;
200 }
201 b'\\' if body.get(index + 1) != Some(&b'u') => index += 2,
204 b'\\' => {
205 let Some(unit) = hex_escape_at(body, index) else {
206 index += 2;
207 continue;
208 };
209
210 let is_pair = (0xD800..0xDC00).contains(&unit)
211 && body.get(index + 6) == Some(&b'\\')
212 && body.get(index + 7) == Some(&b'u')
213 && hex_escape_at(body, index + 6)
214 .is_some_and(|low| (0xDC00..0xE000).contains(&low));
215
216 if is_pair {
217 index += 12;
218 continue;
219 }
220
221 if (0xD800..0xE000).contains(&unit) {
222 let output = repaired.get_or_insert_with(|| body.to_vec());
223 output[index..index + 6].copy_from_slice(br"\ufffd");
224 }
225 index += 6;
226 }
227 _ => index += 1,
228 }
229 }
230
231 repaired
232}
233
234struct WriteCommand {
243 frame: Vec<u8>,
244 ack: oneshot::Sender<Result<(), std::io::Error>>,
245}
246
247pub struct JsonRpcClient {
258 request_id: AtomicU64,
259 write_tx: mpsc::UnboundedSender<WriteCommand>,
266 pending_requests: Arc<RwLock<HashMap<u64, PendingRequest>>>,
267 notification_tx: broadcast::Sender<JsonRpcNotification>,
268 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
269 read_task: Mutex<Option<JoinHandle<()>>>,
270 write_task: Mutex<Option<JoinHandle<()>>>,
271}
272
273impl JsonRpcClient {
274 pub fn new(
281 writer: impl AsyncWrite + Unpin + Send + 'static,
282 reader: impl AsyncRead + Unpin + Send + 'static,
283 notification_tx: broadcast::Sender<JsonRpcNotification>,
284 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
285 ) -> Self {
286 let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteCommand>();
287
288 let writer_span = tracing::error_span!("jsonrpc_write_loop");
289 let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span));
290
291 let client = Self {
292 request_id: AtomicU64::new(1),
293 write_tx,
294 pending_requests: Arc::new(RwLock::new(HashMap::new())),
295 notification_tx,
296 request_tx,
297 read_task: Mutex::new(None),
298 write_task: Mutex::new(Some(write_task)),
299 };
300
301 let pending_requests = client.pending_requests.clone();
302 let notification_tx_clone = client.notification_tx.clone();
303 let request_tx_clone = client.request_tx.clone();
304 let reader_span = tracing::error_span!("jsonrpc_read_loop");
305
306 let read_task = tokio::spawn(
307 async move {
308 Self::read_loop(
309 reader,
310 pending_requests,
311 notification_tx_clone,
312 request_tx_clone,
313 )
314 .await;
315 }
316 .instrument(reader_span),
317 );
318 *client.read_task.lock() = Some(read_task);
319
320 client
321 }
322
323 pub(crate) fn force_close(&self) {
324 if let Some(task) = self.read_task.lock().take() {
325 task.abort();
326 }
327 if let Some(task) = self.write_task.lock().take() {
328 task.abort();
329 }
330 self.pending_requests.write().clear();
331 }
332
333 async fn write_loop(
346 mut writer: impl AsyncWrite + Unpin + Send + 'static,
347 mut rx: mpsc::UnboundedReceiver<WriteCommand>,
348 ) {
349 while let Some(WriteCommand { frame, ack }) = rx.recv().await {
350 let result = async {
351 writer.write_all(&frame).await?;
352 writer.flush().await?;
353 Ok::<_, std::io::Error>(())
354 }
355 .await;
356
357 let _ = ack.send(result);
361 }
362 }
363
364 async fn read_loop(
365 reader: impl AsyncRead + Unpin + Send,
366 pending_requests: Arc<RwLock<HashMap<u64, PendingRequest>>>,
367 notification_tx: broadcast::Sender<JsonRpcNotification>,
368 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
369 ) {
370 let mut reader = BufReader::new(reader);
371
372 loop {
373 match Self::read_message(&mut reader).await {
374 Ok(Some(message)) => match message {
375 JsonRpcMessage::Response(mut response) => {
376 let id = response.id;
377 let pending = pending_requests.write().remove(&id);
378 if let Some(PendingRequest {
379 sender,
380 inline_callback,
381 }) = pending
382 {
383 if let Some(cb) = inline_callback
389 && response.error.is_none()
390 {
391 let cb_outcome =
392 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
393 cb(&response)
394 }));
395 match cb_outcome {
396 Ok(Ok(())) => {}
397 Ok(Err(error)) => {
398 response.result = None;
399 response.error = Some(JsonRpcError {
400 code: -32603,
401 message: error.to_string(),
402 data: None,
403 });
404 }
405 Err(panic) => {
406 let message = panic
407 .downcast_ref::<&'static str>()
408 .map(|s| (*s).to_string())
409 .or_else(|| panic.downcast_ref::<String>().cloned())
410 .unwrap_or_else(|| {
411 "inline response callback panicked".to_string()
412 });
413 response.result = None;
414 response.error = Some(JsonRpcError {
415 code: -32603,
416 message,
417 data: None,
418 });
419 }
420 }
421 }
422 if sender.send(response).is_err() {
423 warn!(request_id = %id, "failed to send response for request");
424 }
425 } else {
426 warn!(request_id = %id, "received response for unknown request id");
427 }
428 }
429 JsonRpcMessage::Notification(notification) => {
430 let _ = notification_tx.send(notification);
431 }
432 JsonRpcMessage::Request(request) => {
433 if request_tx.send(request).is_err() {
434 warn!("failed to forward JSON-RPC request, channel closed");
435 }
436 }
437 },
438 Ok(None) => {
439 break;
440 }
441 Err(e) => {
442 error!(error = %e, "error reading from CLI");
443 break;
444 }
445 }
446 }
447
448 let mut pending = pending_requests.write();
451 if !pending.is_empty() {
452 warn!(
453 count = pending.len(),
454 "draining pending requests after read loop exit"
455 );
456 pending.clear();
457 }
458 }
459
460 async fn read_message(
461 reader: &mut BufReader<impl AsyncRead + Unpin>,
462 ) -> Result<Option<JsonRpcMessage>, Error> {
463 let mut line = String::new();
464 let mut content_length = None;
465
466 loop {
467 line.clear();
468 if reader.read_line(&mut line).await? == 0 {
469 return Ok(None);
470 }
471
472 let trimmed = line.trim();
473 if trimmed.is_empty() {
474 break;
475 }
476
477 if let Some(value) = trimmed.strip_prefix(CONTENT_LENGTH_HEADER) {
478 content_length = Some(value.trim().parse::<usize>().map_err(|_| {
479 Error::from(ErrorKind::Protocol(
480 ProtocolErrorKind::InvalidContentLength(value.trim().to_string()),
481 ))
482 })?);
483 }
484 }
485
486 let Some(length) = content_length else {
487 return Err(ErrorKind::Protocol(ProtocolErrorKind::MissingContentLength).into());
488 };
489
490 let mut body = vec![0u8; length];
491 reader.read_exact(&mut body).await?;
492
493 match serde_json::from_slice::<JsonRpcMessage>(&body) {
494 Ok(message) => Ok(Some(message)),
495 Err(error) => {
496 match repair_lone_surrogates(&body)
499 .and_then(|repaired| serde_json::from_slice::<JsonRpcMessage>(&repaired).ok())
500 {
501 Some(message) => {
502 warn!(
503 error = %error,
504 length,
505 "recovered JSON-RPC frame containing unpaired UTF-16 surrogates"
506 );
507 Ok(Some(message))
508 }
509 None => Err(error.into()),
510 }
511 }
512 }
513 }
514
515 #[allow(dead_code, reason = "public API exported via crate::JsonRpcClient")]
526 pub async fn send_request(
527 &self,
528 method: &str,
529 params: Option<serde_json::Value>,
530 ) -> Result<JsonRpcResponse, Error> {
531 self.send_request_with_inline_callback(method, params, None)
532 .await
533 }
534
535 pub(crate) async fn send_request_with_inline_callback(
552 &self,
553 method: &str,
554 params: Option<serde_json::Value>,
555 inline_callback: Option<InlineResponseCallback>,
556 ) -> Result<JsonRpcResponse, Error> {
557 let request_start = Instant::now();
558 let id = self.request_id.fetch_add(1, Ordering::SeqCst);
559 let request = JsonRpcRequest::new(id, method, params);
560
561 let (tx, rx) = oneshot::channel();
562 self.pending_requests.write().insert(
563 id,
564 PendingRequest {
565 sender: tx,
566 inline_callback,
567 },
568 );
569
570 let mut guard = PendingGuard {
575 map: &self.pending_requests,
576 id,
577 armed: true,
578 };
579
580 if let Err(error) = self.write(&request).await {
584 warn!(
585 elapsed_ms = request_start.elapsed().as_millis(),
586 method = %method,
587 request_id = id,
588 status = "failed",
589 error = %error,
590 "JsonRpcClient::send_request JSON-RPC request finished"
591 );
592 return Err(error);
593 }
594
595 let response = match rx.await {
596 Ok(response) => response,
597 Err(_) => {
598 let error = ErrorKind::Protocol(ProtocolErrorKind::RequestCancelled).into();
599 warn!(
600 elapsed_ms = request_start.elapsed().as_millis(),
601 method = %method,
602 request_id = id,
603 status = "failed",
604 error = %error,
605 "JsonRpcClient::send_request JSON-RPC request finished"
606 );
607 return Err(error);
608 }
609 };
610 guard.disarm();
611 if let Some(error) = &response.error {
612 warn!(
613 elapsed_ms = request_start.elapsed().as_millis(),
614 method = %method,
615 request_id = id,
616 status = "failed",
617 code = error.code,
618 error = %error.message,
619 "JsonRpcClient::send_request JSON-RPC request finished"
620 );
621 } else {
622 debug!(
623 elapsed_ms = request_start.elapsed().as_millis(),
624 method = %method,
625 request_id = id,
626 status = "succeeded",
627 "JsonRpcClient::send_request JSON-RPC request finished"
628 );
629 }
630 Ok(response)
631 }
632
633 pub async fn write<T: serde::Serialize>(&self, message: &T) -> Result<(), Error> {
642 let body = serde_json::to_vec(message)?;
643 let mut frame = Vec::with_capacity(CONTENT_LENGTH_HEADER.len() + 16 + body.len() + 4);
644 frame.extend_from_slice(CONTENT_LENGTH_HEADER.as_bytes());
645 frame.extend_from_slice(body.len().to_string().as_bytes());
646 frame.extend_from_slice(b"\r\n\r\n");
647 frame.extend_from_slice(&body);
648
649 let (ack_tx, ack_rx) = oneshot::channel();
650 self.write_tx
651 .send(WriteCommand { frame, ack: ack_tx })
652 .map_err(|_| {
653 Error::from(std::io::Error::new(
654 std::io::ErrorKind::BrokenPipe,
655 "writer actor has shut down",
656 ))
657 })?;
658
659 match ack_rx.await {
660 Ok(Ok(())) => Ok(()),
661 Ok(Err(e)) => Err(Error::from(e)),
662 Err(_) => Err(Error::from(std::io::Error::new(
663 std::io::ErrorKind::BrokenPipe,
664 "writer actor dropped ack without responding",
665 ))),
666 }
667 }
668}
669
670struct PendingGuard<'a> {
674 map: &'a RwLock<HashMap<u64, PendingRequest>>,
675 id: u64,
676 armed: bool,
677}
678
679impl PendingGuard<'_> {
680 fn disarm(&mut self) {
681 self.armed = false;
682 }
683}
684
685impl Drop for PendingGuard<'_> {
686 fn drop(&mut self) {
687 if self.armed {
688 self.map.write().remove(&self.id);
689 }
690 }
691}
692
693#[cfg(test)]
694mod tests {
695 use super::*;
696
697 #[test]
698 fn deserialize_notification() {
699 let json = r#"{"jsonrpc":"2.0","method":"session.event","params":{"id":"e1"}}"#;
700 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
701 assert!(matches!(msg, JsonRpcMessage::Notification(n) if n.method == "session.event"));
702 }
703
704 #[test]
705 fn deserialize_request() {
706 let json =
707 r#"{"jsonrpc":"2.0","id":5,"method":"permission.request","params":{"kind":"shell"}}"#;
708 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
709 assert!(
710 matches!(msg, JsonRpcMessage::Request(r) if r.id == 5 && r.method == "permission.request")
711 );
712 }
713
714 #[test]
715 fn deserialize_response_with_result() {
716 let json = r#"{"jsonrpc":"2.0","id":3,"result":{"ok":true}}"#;
717 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
718 assert!(matches!(msg, JsonRpcMessage::Response(r) if r.id == 3 && !r.is_error()));
719 }
720
721 #[test]
722 fn deserialize_error_response() {
723 let json =
724 r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32600,"message":"Invalid Request"}}"#;
725 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
726 match msg {
727 JsonRpcMessage::Response(r) => {
728 assert!(r.is_error());
729 let err = r.error.unwrap();
730 assert_eq!(err.code, -32600);
731 assert_eq!(err.message, "Invalid Request");
732 }
733 other => panic!("expected Response, got {other:?}"),
734 }
735 }
736
737 #[test]
738 fn deserialize_rejects_non_object() {
739 let result = serde_json::from_str::<JsonRpcMessage>(r#""not an object""#);
740 assert!(result.is_err());
741 }
742
743 #[test]
744 fn request_new_sets_version() {
745 let req = JsonRpcRequest::new(42, "test.method", None);
746 assert_eq!(req.jsonrpc, "2.0");
747 assert_eq!(req.id, 42);
748 assert_eq!(req.method, "test.method");
749 assert!(req.params.is_none());
750 }
751
752 #[test]
753 fn request_serializes_camel_case() {
754 let req = JsonRpcRequest::new(1, "ping", Some(serde_json::json!({})));
755 let json = serde_json::to_string(&req).unwrap();
756 assert!(json.contains(r#""jsonrpc":"2.0""#));
757 assert!(json.contains(r#""id":1"#));
758 assert!(json.contains(r#""method":"ping""#));
759 }
760
761 #[test]
762 fn notification_without_params_omits_field() {
763 let n = JsonRpcNotification {
764 jsonrpc: "2.0".into(),
765 method: "ping".into(),
766 params: None,
767 };
768 let json = serde_json::to_string(&n).unwrap();
769 assert!(!json.contains("params"));
770 }
771
772 #[test]
773 fn response_without_error_omits_field() {
774 let r = JsonRpcResponse {
775 jsonrpc: "2.0".into(),
776 id: 1,
777 result: Some(serde_json::json!(true)),
778 error: None,
779 };
780 let json = serde_json::to_string(&r).unwrap();
781 assert!(!json.contains("error"));
782 }
783}