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 value = Value::deserialize(deserializer)?;
128 let obj = value
129 .as_object()
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 if has_id && has_method {
136 JsonRpcRequest::deserialize(value)
137 .map(JsonRpcMessage::Request)
138 .map_err(serde::de::Error::custom)
139 } else if has_id {
140 JsonRpcResponse::deserialize(value)
141 .map(JsonRpcMessage::Response)
142 .map_err(serde::de::Error::custom)
143 } else {
144 JsonRpcNotification::deserialize(value)
145 .map(JsonRpcMessage::Notification)
146 .map_err(serde::de::Error::custom)
147 }
148 }
149}
150
151impl JsonRpcRequest {
152 pub fn new(id: u64, method: &str, params: Option<Value>) -> Self {
154 Self {
155 jsonrpc: "2.0".to_string(),
156 id,
157 method: method.to_string(),
158 params,
159 }
160 }
161}
162
163impl JsonRpcResponse {
164 #[allow(dead_code)]
166 pub fn is_error(&self) -> bool {
167 self.error.is_some()
168 }
169}
170
171const CONTENT_LENGTH_HEADER: &str = "Content-Length: ";
172
173fn repair_lone_surrogates(body: &[u8]) -> Option<Vec<u8>> {
178 fn hex_escape_at(body: &[u8], index: usize) -> Option<u16> {
179 let digits = body.get(index + 2..index + 6)?;
180 let text = std::str::from_utf8(digits).ok()?;
181 u16::from_str_radix(text, 16).ok()
182 }
183
184 let mut repaired = None;
185 let mut in_string = false;
186 let mut index = 0;
187
188 while index < body.len() {
189 let byte = body[index];
190
191 if !in_string {
192 in_string = byte == b'"';
193 index += 1;
194 continue;
195 }
196
197 match byte {
198 b'"' => {
199 in_string = false;
200 index += 1;
201 }
202 b'\\' if body.get(index + 1) != Some(&b'u') => index += 2,
205 b'\\' => {
206 let Some(unit) = hex_escape_at(body, index) else {
207 index += 2;
208 continue;
209 };
210
211 let is_pair = (0xD800..0xDC00).contains(&unit)
212 && body.get(index + 6) == Some(&b'\\')
213 && body.get(index + 7) == Some(&b'u')
214 && hex_escape_at(body, index + 6)
215 .is_some_and(|low| (0xDC00..0xE000).contains(&low));
216
217 if is_pair {
218 index += 12;
219 continue;
220 }
221
222 if (0xD800..0xE000).contains(&unit) {
223 let output = repaired.get_or_insert_with(|| body.to_vec());
224 output[index..index + 6].copy_from_slice(br"\ufffd");
225 }
226 index += 6;
227 }
228 _ => index += 1,
229 }
230 }
231
232 repaired
233}
234
235struct WriteCommand {
244 frame: Vec<u8>,
245 ack: oneshot::Sender<Result<(), std::io::Error>>,
246}
247
248pub struct JsonRpcClient {
259 request_id: AtomicU64,
260 write_tx: mpsc::UnboundedSender<WriteCommand>,
267 pending_requests: Arc<RwLock<HashMap<u64, PendingRequest>>>,
268 notification_tx: broadcast::Sender<JsonRpcNotification>,
269 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
270 connection_closed: CancellationToken,
271 read_task: Mutex<Option<JoinHandle<()>>>,
272 write_task: Mutex<Option<JoinHandle<()>>>,
273}
274
275impl JsonRpcClient {
276 pub fn new(
283 writer: impl AsyncWrite + Unpin + Send + 'static,
284 reader: impl AsyncRead + Unpin + Send + 'static,
285 notification_tx: broadcast::Sender<JsonRpcNotification>,
286 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
287 ) -> Self {
288 let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteCommand>();
289
290 let writer_span = tracing::error_span!("jsonrpc_write_loop");
291 let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span));
292
293 let client = Self {
294 request_id: AtomicU64::new(1),
295 write_tx,
296 pending_requests: Arc::new(RwLock::new(HashMap::new())),
297 notification_tx,
298 request_tx,
299 connection_closed: CancellationToken::new(),
300 read_task: Mutex::new(None),
301 write_task: Mutex::new(Some(write_task)),
302 };
303
304 let pending_requests = client.pending_requests.clone();
305 let notification_tx_clone = client.notification_tx.clone();
306 let request_tx_clone = client.request_tx.clone();
307 let connection_closed = client.connection_closed.clone();
308 let reader_span = tracing::error_span!("jsonrpc_read_loop");
309
310 let read_task = tokio::spawn(
311 async move {
312 Self::read_loop(
313 reader,
314 pending_requests,
315 notification_tx_clone,
316 request_tx_clone,
317 )
318 .await;
319 connection_closed.cancel();
320 }
321 .instrument(reader_span),
322 );
323 *client.read_task.lock() = Some(read_task);
324
325 client
326 }
327
328 pub(crate) fn force_close(&self) {
329 self.connection_closed.cancel();
330 if let Some(task) = self.read_task.lock().take() {
331 task.abort();
332 }
333 if let Some(task) = self.write_task.lock().take() {
334 task.abort();
335 }
336 self.pending_requests.write().clear();
337 }
338
339 pub(crate) fn connection_closed_token(&self) -> CancellationToken {
340 self.connection_closed.child_token()
341 }
342
343 async fn write_loop(
356 mut writer: impl AsyncWrite + Unpin + Send + 'static,
357 mut rx: mpsc::UnboundedReceiver<WriteCommand>,
358 ) {
359 while let Some(WriteCommand { frame, ack }) = rx.recv().await {
360 let result = async {
361 writer.write_all(&frame).await?;
362 writer.flush().await?;
363 Ok::<_, std::io::Error>(())
364 }
365 .await;
366
367 let _ = ack.send(result);
371 }
372 }
373
374 async fn read_loop(
375 reader: impl AsyncRead + Unpin + Send,
376 pending_requests: Arc<RwLock<HashMap<u64, PendingRequest>>>,
377 notification_tx: broadcast::Sender<JsonRpcNotification>,
378 request_tx: mpsc::UnboundedSender<JsonRpcRequest>,
379 ) {
380 let mut reader = BufReader::new(reader);
381
382 loop {
383 match Self::read_message(&mut reader).await {
384 Ok(Some(message)) => match message {
385 JsonRpcMessage::Response(mut response) => {
386 let id = response.id;
387 let pending = pending_requests.write().remove(&id);
388 if let Some(PendingRequest {
389 sender,
390 inline_callback,
391 }) = pending
392 {
393 if let Some(cb) = inline_callback
399 && response.error.is_none()
400 {
401 let cb_outcome =
402 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
403 cb(&response)
404 }));
405 match cb_outcome {
406 Ok(Ok(())) => {}
407 Ok(Err(error)) => {
408 response.result = None;
409 response.error = Some(JsonRpcError {
410 code: -32603,
411 message: error.to_string(),
412 data: None,
413 });
414 }
415 Err(panic) => {
416 let message = panic
417 .downcast_ref::<&'static str>()
418 .map(|s| (*s).to_string())
419 .or_else(|| panic.downcast_ref::<String>().cloned())
420 .unwrap_or_else(|| {
421 "inline response callback panicked".to_string()
422 });
423 response.result = None;
424 response.error = Some(JsonRpcError {
425 code: -32603,
426 message,
427 data: None,
428 });
429 }
430 }
431 }
432 if sender.send(response).is_err() {
433 warn!(request_id = %id, "failed to send response for request");
434 }
435 } else {
436 warn!(request_id = %id, "received response for unknown request id");
437 }
438 }
439 JsonRpcMessage::Notification(notification) => {
440 let _ = notification_tx.send(notification);
441 }
442 JsonRpcMessage::Request(request) => {
443 if request_tx.send(request).is_err() {
444 warn!("failed to forward JSON-RPC request, channel closed");
445 }
446 }
447 },
448 Ok(None) => {
449 break;
450 }
451 Err(e) => {
452 error!(error = %e, "error reading from CLI");
453 break;
454 }
455 }
456 }
457
458 let mut pending = pending_requests.write();
461 if !pending.is_empty() {
462 warn!(
463 count = pending.len(),
464 "draining pending requests after read loop exit"
465 );
466 pending.clear();
467 }
468 }
469
470 async fn read_message(
471 reader: &mut BufReader<impl AsyncRead + Unpin>,
472 ) -> Result<Option<JsonRpcMessage>, Error> {
473 let mut line = String::new();
474 let mut content_length = None;
475
476 loop {
477 line.clear();
478 if reader.read_line(&mut line).await? == 0 {
479 return Ok(None);
480 }
481
482 let trimmed = line.trim();
483 if trimmed.is_empty() {
484 break;
485 }
486
487 if let Some(value) = trimmed.strip_prefix(CONTENT_LENGTH_HEADER) {
488 content_length = Some(value.trim().parse::<usize>().map_err(|_| {
489 Error::from(ErrorKind::Protocol(
490 ProtocolErrorKind::InvalidContentLength(value.trim().to_string()),
491 ))
492 })?);
493 }
494 }
495
496 let Some(length) = content_length else {
497 return Err(ErrorKind::Protocol(ProtocolErrorKind::MissingContentLength).into());
498 };
499
500 let mut body = vec![0u8; length];
501 reader.read_exact(&mut body).await?;
502
503 match serde_json::from_slice::<JsonRpcMessage>(&body) {
504 Ok(message) => Ok(Some(message)),
505 Err(error) => {
506 match repair_lone_surrogates(&body)
509 .and_then(|repaired| serde_json::from_slice::<JsonRpcMessage>(&repaired).ok())
510 {
511 Some(message) => {
512 warn!(
513 error = %error,
514 length,
515 "recovered JSON-RPC frame containing unpaired UTF-16 surrogates"
516 );
517 Ok(Some(message))
518 }
519 None => Err(error.into()),
520 }
521 }
522 }
523 }
524
525 #[allow(dead_code, reason = "public API exported via crate::JsonRpcClient")]
536 pub async fn send_request(
537 &self,
538 method: &str,
539 params: Option<serde_json::Value>,
540 ) -> Result<JsonRpcResponse, Error> {
541 self.send_request_with_inline_callback(method, params, None)
542 .await
543 }
544
545 pub(crate) async fn send_request_with_inline_callback(
562 &self,
563 method: &str,
564 params: Option<serde_json::Value>,
565 inline_callback: Option<InlineResponseCallback>,
566 ) -> Result<JsonRpcResponse, Error> {
567 let request_start = Instant::now();
568 let id = self.request_id.fetch_add(1, Ordering::SeqCst);
569 let request = JsonRpcRequest::new(id, method, params);
570
571 let (tx, rx) = oneshot::channel();
572 self.pending_requests.write().insert(
573 id,
574 PendingRequest {
575 sender: tx,
576 inline_callback,
577 },
578 );
579
580 let mut guard = PendingGuard {
585 map: &self.pending_requests,
586 id,
587 armed: true,
588 };
589
590 if let Err(error) = self.write(&request).await {
594 warn!(
595 elapsed_ms = request_start.elapsed().as_millis(),
596 method = %method,
597 request_id = id,
598 status = "failed",
599 error = %error,
600 "JsonRpcClient::send_request JSON-RPC request finished"
601 );
602 return Err(error);
603 }
604
605 let response = match rx.await {
606 Ok(response) => response,
607 Err(_) => {
608 let error = ErrorKind::Protocol(ProtocolErrorKind::RequestCancelled).into();
609 warn!(
610 elapsed_ms = request_start.elapsed().as_millis(),
611 method = %method,
612 request_id = id,
613 status = "failed",
614 error = %error,
615 "JsonRpcClient::send_request JSON-RPC request finished"
616 );
617 return Err(error);
618 }
619 };
620 guard.disarm();
621 if let Some(error) = &response.error {
622 warn!(
623 elapsed_ms = request_start.elapsed().as_millis(),
624 method = %method,
625 request_id = id,
626 status = "failed",
627 code = error.code,
628 error = %error.message,
629 "JsonRpcClient::send_request JSON-RPC request finished"
630 );
631 } else {
632 debug!(
633 elapsed_ms = request_start.elapsed().as_millis(),
634 method = %method,
635 request_id = id,
636 status = "succeeded",
637 "JsonRpcClient::send_request JSON-RPC request finished"
638 );
639 }
640 Ok(response)
641 }
642
643 pub async fn write<T: serde::Serialize>(&self, message: &T) -> Result<(), Error> {
652 let body = serde_json::to_vec(message)?;
653 let mut frame = Vec::with_capacity(CONTENT_LENGTH_HEADER.len() + 16 + body.len() + 4);
654 frame.extend_from_slice(CONTENT_LENGTH_HEADER.as_bytes());
655 frame.extend_from_slice(body.len().to_string().as_bytes());
656 frame.extend_from_slice(b"\r\n\r\n");
657 frame.extend_from_slice(&body);
658
659 let (ack_tx, ack_rx) = oneshot::channel();
660 self.write_tx
661 .send(WriteCommand { frame, ack: ack_tx })
662 .map_err(|_| {
663 Error::from(std::io::Error::new(
664 std::io::ErrorKind::BrokenPipe,
665 "writer actor has shut down",
666 ))
667 })?;
668
669 match ack_rx.await {
670 Ok(Ok(())) => Ok(()),
671 Ok(Err(e)) => Err(Error::from(e)),
672 Err(_) => Err(Error::from(std::io::Error::new(
673 std::io::ErrorKind::BrokenPipe,
674 "writer actor dropped ack without responding",
675 ))),
676 }
677 }
678}
679
680struct PendingGuard<'a> {
684 map: &'a RwLock<HashMap<u64, PendingRequest>>,
685 id: u64,
686 armed: bool,
687}
688
689impl PendingGuard<'_> {
690 fn disarm(&mut self) {
691 self.armed = false;
692 }
693}
694
695impl Drop for PendingGuard<'_> {
696 fn drop(&mut self) {
697 if self.armed {
698 self.map.write().remove(&self.id);
699 }
700 }
701}
702
703#[cfg(test)]
704mod tests {
705 use super::*;
706
707 #[test]
708 fn deserialize_notification() {
709 let json = r#"{"jsonrpc":"2.0","method":"session.event","params":{"id":"e1"}}"#;
710 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
711 assert!(matches!(msg, JsonRpcMessage::Notification(n) if n.method == "session.event"));
712 }
713
714 #[test]
715 fn deserialize_request() {
716 let json =
717 r#"{"jsonrpc":"2.0","id":5,"method":"permission.request","params":{"kind":"shell"}}"#;
718 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
719 assert!(
720 matches!(msg, JsonRpcMessage::Request(r) if r.id == 5 && r.method == "permission.request")
721 );
722 }
723
724 #[test]
725 fn deserialize_response_with_result() {
726 let json = r#"{"jsonrpc":"2.0","id":3,"result":{"ok":true}}"#;
727 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
728 assert!(matches!(msg, JsonRpcMessage::Response(r) if r.id == 3 && !r.is_error()));
729 }
730
731 #[test]
732 fn deserialize_error_response() {
733 let json =
734 r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32600,"message":"Invalid Request"}}"#;
735 let msg: JsonRpcMessage = serde_json::from_str(json).unwrap();
736 match msg {
737 JsonRpcMessage::Response(r) => {
738 assert!(r.is_error());
739 let err = r.error.unwrap();
740 assert_eq!(err.code, -32600);
741 assert_eq!(err.message, "Invalid Request");
742 }
743 other => panic!("expected Response, got {other:?}"),
744 }
745 }
746
747 #[test]
748 fn deserialize_rejects_non_object() {
749 let result = serde_json::from_str::<JsonRpcMessage>(r#""not an object""#);
750 assert!(result.is_err());
751 }
752
753 #[test]
754 fn request_new_sets_version() {
755 let req = JsonRpcRequest::new(42, "test.method", None);
756 assert_eq!(req.jsonrpc, "2.0");
757 assert_eq!(req.id, 42);
758 assert_eq!(req.method, "test.method");
759 assert!(req.params.is_none());
760 }
761
762 #[test]
763 fn request_serializes_camel_case() {
764 let req = JsonRpcRequest::new(1, "ping", Some(serde_json::json!({})));
765 let json = serde_json::to_string(&req).unwrap();
766 assert!(json.contains(r#""jsonrpc":"2.0""#));
767 assert!(json.contains(r#""id":1"#));
768 assert!(json.contains(r#""method":"ping""#));
769 }
770
771 #[test]
772 fn notification_without_params_omits_field() {
773 let n = JsonRpcNotification {
774 jsonrpc: "2.0".into(),
775 method: "ping".into(),
776 params: None,
777 };
778 let json = serde_json::to_string(&n).unwrap();
779 assert!(!json.contains("params"));
780 }
781
782 #[test]
783 fn response_without_error_omits_field() {
784 let r = JsonRpcResponse {
785 jsonrpc: "2.0".into(),
786 id: 1,
787 result: Some(serde_json::json!(true)),
788 error: None,
789 };
790 let json = serde_json::to_string(&r).unwrap();
791 assert!(!json.contains("error"));
792 }
793}