use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; use parking_lot::{Mutex, RwLock}; use serde::{Deserialize, Serialize}; use serde_json::Value; use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader}; use tokio::sync::{broadcast, mpsc, oneshot}; use tokio::task::JoinHandle; use tracing::{Instrument, debug, error, warn}; use crate::{Error, ErrorKind, ProtocolErrorKind}; /// Callback invoked synchronously by the JSON-RPC read loop the instant a /// successful response is parsed, before the response is delivered to the /// awaiter and before the read loop dispatches the next message. Use this /// when client-side state (for example, registering a server-assigned /// session id with the router) must be visible to any subsequent /// notification on the same connection. /// /// If the callback returns an error, that error is delivered to the /// awaiter in place of the response. pub(crate) type InlineResponseCallback = Box Result<(), Error> + Send + Sync>; /// Internal pairing of the response delivery channel with an optional /// inline callback that the read loop runs synchronously before delivery. struct PendingRequest { sender: oneshot::Sender, inline_callback: Option, } /// A JSON-RPC 2.0 request message. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct JsonRpcRequest { /// Protocol version (always `"2.0"`). pub jsonrpc: String, /// Request ID for correlating responses. pub id: u64, /// RPC method name. pub method: String, /// Optional method parameters. #[serde(skip_serializing_if = "Option::is_none")] pub params: Option, } /// A JSON-RPC 2.0 response message. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct JsonRpcResponse { /// Protocol version (always `"2.0"`). pub jsonrpc: String, /// Request ID this response correlates to. pub id: u64, /// Success payload (mutually exclusive with `error`). #[serde(skip_serializing_if = "Option::is_none")] pub result: Option, /// Error payload (mutually exclusive with `result`). #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, } /// A JSON-RPC 2.0 error object. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct JsonRpcError { /// Numeric error code. pub code: i32, /// Human-readable error description. pub message: String, /// Optional structured error data. #[serde(skip_serializing_if = "Option::is_none")] pub data: Option, } /// Standard JSON-RPC 2.0 error codes. pub mod error_codes { /// Method not found (-32601). pub const METHOD_NOT_FOUND: i32 = -32601; /// Invalid method parameters (-32602). pub const INVALID_PARAMS: i32 = -32602; /// Internal server error (-32603). #[allow(dead_code, reason = "standard JSON-RPC code, reserved for future use")] pub const INTERNAL_ERROR: i32 = -32603; } /// A JSON-RPC 2.0 notification (no `id`, no response expected). #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct JsonRpcNotification { /// Protocol version (always `"2.0"`). pub jsonrpc: String, /// Notification method name. pub method: String, /// Optional notification parameters. #[serde(skip_serializing_if = "Option::is_none")] pub params: Option, } /// A parsed JSON-RPC 2.0 message — request, response, or notification. #[derive(Debug, Clone, Serialize)] pub enum JsonRpcMessage { /// An incoming or outgoing request. Request(JsonRpcRequest), /// A response to a previous request. Response(JsonRpcResponse), /// A fire-and-forget notification. Notification(JsonRpcNotification), } /// Custom deserializer that dispatches based on field presence instead of /// `#[serde(untagged)]` which tries each variant sequentially (3× parse /// attempts for Notification — the hot-path streaming variant). /// /// Dispatch logic: /// - has `id` + has `method` → Request /// - has `id` + no `method` → Response /// - no `id` → Notification impl<'de> Deserialize<'de> for JsonRpcMessage { fn deserialize(deserializer: D) -> Result where D: serde::Deserializer<'de>, { let value = Value::deserialize(deserializer)?; let obj = value .as_object() .ok_or_else(|| serde::de::Error::custom("expected a JSON object"))?; let has_id = obj.contains_key("id"); let has_method = obj.contains_key("method"); if has_id && has_method { JsonRpcRequest::deserialize(value) .map(JsonRpcMessage::Request) .map_err(serde::de::Error::custom) } else if has_id { JsonRpcResponse::deserialize(value) .map(JsonRpcMessage::Response) .map_err(serde::de::Error::custom) } else { JsonRpcNotification::deserialize(value) .map(JsonRpcMessage::Notification) .map_err(serde::de::Error::custom) } } } impl JsonRpcRequest { /// Create a new JSON-RPC request with the given ID, method, and params. pub fn new(id: u64, method: &str, params: Option) -> Self { Self { jsonrpc: "2.0".to_string(), id, method: method.to_string(), params, } } } impl JsonRpcResponse { /// Returns `true` if this response contains an error. #[allow(dead_code)] pub fn is_error(&self) -> bool { self.error.is_some() } } const CONTENT_LENGTH_HEADER: &str = "Content-Length: "; /// Rewrites unpaired UTF-16 surrogate escapes to `\uFFFD`. /// /// Returns `None` when the body contains no unpaired surrogate, so valid /// frames do not incur a repair allocation. fn repair_lone_surrogates(body: &[u8]) -> Option> { fn hex_escape_at(body: &[u8], index: usize) -> Option { let digits = body.get(index + 2..index + 6)?; let text = std::str::from_utf8(digits).ok()?; u16::from_str_radix(text, 16).ok() } let mut repaired = None; let mut in_string = false; let mut index = 0; while index < body.len() { let byte = body[index]; if !in_string { in_string = byte == b'"'; index += 1; continue; } match byte { b'"' => { in_string = false; index += 1; } // Consume non-Unicode escapes whole so an escaped backslash cannot // be mistaken for the start of a surrogate escape. b'\\' if body.get(index + 1) != Some(&b'u') => index += 2, b'\\' => { let Some(unit) = hex_escape_at(body, index) else { index += 2; continue; }; let is_pair = (0xD800..0xDC00).contains(&unit) && body.get(index + 6) == Some(&b'\\') && body.get(index + 7) == Some(&b'u') && hex_escape_at(body, index + 6) .is_some_and(|low| (0xDC00..0xE000).contains(&low)); if is_pair { index += 12; continue; } if (0xD800..0xE000).contains(&unit) { let output = repaired.get_or_insert_with(|| body.to_vec()); output[index..index + 6].copy_from_slice(br"\ufffd"); } index += 6; } _ => index += 1, } } repaired } /// One framed JSON-RPC message handed to the writer actor. /// /// `frame` is the fully serialized bytes (header + body); the caller pays /// the serde cost synchronously before enqueueing so the actor never sees a /// `Result` from JSON encoding. `ack` resolves once the bytes have been /// fully written and flushed (or the underlying I/O reports an error). If /// the caller drops the `oneshot::Receiver`, the actor still completes the /// frame — caller cancellation cannot desync the wire. struct WriteCommand { frame: Vec, ack: oneshot::Sender>, } /// Low-level JSON-RPC 2.0 client over Content-Length-framed streams. /// /// # Cancel safety /// /// All public methods (`write`, `send_request`) are **cancel-safe**: the /// actual bytes hit the wire on a dedicated background actor task, so /// dropping the caller's future after `await` returns `Pending` cannot /// produce a partial frame on the wire. Frames either land atomically or /// the underlying I/O fails. See `cancel-safety review` artifact for the /// full RFD-400 reasoning. pub struct JsonRpcClient { request_id: AtomicU64, /// Sender side of the writer actor's command queue. Public methods /// pre-serialize their frames and enqueue here; the background actor /// drains the queue and serializes writes onto the underlying /// `AsyncWrite`. Unbounded by design — RFD 400 explicitly permits this /// for cancel-safety, and JSON-RPC frames are small relative to the /// natural request/response back-pressure of the wire. write_tx: mpsc::UnboundedSender, pending_requests: Arc>>, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, read_task: Mutex>>, write_task: Mutex>>, } impl JsonRpcClient { /// Create a new client from async read/write streams. /// /// Spawns two background tasks: a reader that dispatches incoming /// messages to pending request channels, the notification broadcast, /// or the request-forwarding channel; and a writer actor that owns the /// underlying `AsyncWrite` and serializes frames atomically. pub fn new( writer: impl AsyncWrite + Unpin + Send + 'static, reader: impl AsyncRead + Unpin + Send + 'static, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, ) -> Self { let (write_tx, write_rx) = mpsc::unbounded_channel::(); let writer_span = tracing::error_span!("jsonrpc_write_loop"); let write_task = tokio::spawn(Self::write_loop(writer, write_rx).instrument(writer_span)); let client = Self { request_id: AtomicU64::new(1), write_tx, pending_requests: Arc::new(RwLock::new(HashMap::new())), notification_tx, request_tx, read_task: Mutex::new(None), write_task: Mutex::new(Some(write_task)), }; let pending_requests = client.pending_requests.clone(); let notification_tx_clone = client.notification_tx.clone(); let request_tx_clone = client.request_tx.clone(); let reader_span = tracing::error_span!("jsonrpc_read_loop"); let read_task = tokio::spawn( async move { Self::read_loop( reader, pending_requests, notification_tx_clone, request_tx_clone, ) .await; } .instrument(reader_span), ); *client.read_task.lock() = Some(read_task); client } pub(crate) fn force_close(&self) { if let Some(task) = self.read_task.lock().take() { task.abort(); } if let Some(task) = self.write_task.lock().take() { task.abort(); } self.pending_requests.write().clear(); } /// Writer-actor task. Owns the `AsyncWrite`, drains the command queue, /// and writes each frame atomically (header + body + flush) before /// signaling the ack. /// /// Caller-side cancellation cannot interrupt a write in progress: /// dropping the ack `oneshot::Receiver` does not cancel the in-flight /// I/O. Once `WriteCommand` is enqueued the frame is committed to land /// on the wire (or surface an `io::Error` to the ack receiver if the /// transport is broken). /// /// Exits cleanly when all senders drop (channel closes), flushing any /// final buffered bytes. async fn write_loop( mut writer: impl AsyncWrite + Unpin + Send + 'static, mut rx: mpsc::UnboundedReceiver, ) { while let Some(WriteCommand { frame, ack }) = rx.recv().await { let result = async { writer.write_all(&frame).await?; writer.flush().await?; Ok::<_, std::io::Error>(()) } .await; // Caller may have dropped the ack receiver (e.g. their // `await` was cancelled); that's fine — we still completed // the write, which was the whole point. let _ = ack.send(result); } } async fn read_loop( reader: impl AsyncRead + Unpin + Send, pending_requests: Arc>>, notification_tx: broadcast::Sender, request_tx: mpsc::UnboundedSender, ) { let mut reader = BufReader::new(reader); loop { match Self::read_message(&mut reader).await { Ok(Some(message)) => match message { JsonRpcMessage::Response(mut response) => { let id = response.id; let pending = pending_requests.write().remove(&id); if let Some(PendingRequest { sender, inline_callback, }) = pending { // Run the inline callback synchronously on the // read loop so any state it mutates (e.g. // registering a server-assigned session id with // the router) is visible before the loop reads // and dispatches the next message. if let Some(cb) = inline_callback && response.error.is_none() { let cb_outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { cb(&response) })); match cb_outcome { Ok(Ok(())) => {} Ok(Err(error)) => { response.result = None; response.error = Some(JsonRpcError { code: -32603, message: error.to_string(), data: None, }); } Err(panic) => { let message = panic .downcast_ref::<&'static str>() .map(|s| (*s).to_string()) .or_else(|| panic.downcast_ref::().cloned()) .unwrap_or_else(|| { "inline response callback panicked".to_string() }); response.result = None; response.error = Some(JsonRpcError { code: -32603, message, data: None, }); } } } if sender.send(response).is_err() { warn!(request_id = %id, "failed to send response for request"); } } else { warn!(request_id = %id, "received response for unknown request id"); } } JsonRpcMessage::Notification(notification) => { let _ = notification_tx.send(notification); } JsonRpcMessage::Request(request) => { if request_tx.send(request).is_err() { warn!("failed to forward JSON-RPC request, channel closed"); } } }, Ok(None) => { break; } Err(e) => { error!(error = %e, "error reading from CLI"); break; } } } // Drain in-flight requests so callers observe cancellation // instead of hanging on a oneshot receiver. let mut pending = pending_requests.write(); if !pending.is_empty() { warn!( count = pending.len(), "draining pending requests after read loop exit" ); pending.clear(); } } async fn read_message( reader: &mut BufReader, ) -> Result, Error> { let mut line = String::new(); let mut content_length = None; loop { line.clear(); if reader.read_line(&mut line).await? == 0 { return Ok(None); } let trimmed = line.trim(); if trimmed.is_empty() { break; } if let Some(value) = trimmed.strip_prefix(CONTENT_LENGTH_HEADER) { content_length = Some(value.trim().parse::().map_err(|_| { Error::from(ErrorKind::Protocol( ProtocolErrorKind::InvalidContentLength(value.trim().to_string()), )) })?); } } let Some(length) = content_length else { return Err(ErrorKind::Protocol(ProtocolErrorKind::MissingContentLength).into()); }; let mut body = vec![0u8; length]; reader.read_exact(&mut body).await?; match serde_json::from_slice::(&body) { Ok(message) => Ok(Some(message)), Err(error) => { // Dropping an undecodable frame could leave its pending // request waiting forever because this layer has no timeout. match repair_lone_surrogates(&body) .and_then(|repaired| serde_json::from_slice::(&repaired).ok()) { Some(message) => { warn!( error = %error, length, "recovered JSON-RPC frame containing unpaired UTF-16 surrogates" ); Ok(Some(message)) } None => Err(error.into()), } } } } /// Send a JSON-RPC request and wait for the matching response. /// /// # Cancel safety /// /// **Cancel-safe.** The frame is committed to the wire via the writer /// actor before this future yields; cancelling the await drops the /// response oneshot but does not desync the transport. The pending- /// requests map is cleaned up automatically (the `PendingGuard` drop /// removes the entry, and the read loop's response handling tolerates /// a missing entry). #[allow(dead_code, reason = "public API exported via crate::JsonRpcClient")] pub async fn send_request( &self, method: &str, params: Option, ) -> Result { self.send_request_with_inline_callback(method, params, None) .await } /// Send a JSON-RPC request whose response is observed synchronously /// by the read loop *before* it is delivered to the awaiter. /// /// The optional `inline_callback` runs on the JSON-RPC read task the /// instant a successful response is parsed, and before the read loop /// dispatches the next message. This is the only way to perform /// client-side bookkeeping (for example, registering a server- /// assigned session id with the router) that must be visible to any /// notification or request that the server may emit on the same /// connection immediately after the response. /// /// If the callback returns an error or panics, that error is /// surfaced to the awaiter in place of the original response (the /// response payload is discarded and an internal-error JSON-RPC /// error is delivered instead). The error is never propagated back /// to the server and does not crash the read loop. pub(crate) async fn send_request_with_inline_callback( &self, method: &str, params: Option, inline_callback: Option, ) -> Result { let request_start = Instant::now(); let id = self.request_id.fetch_add(1, Ordering::SeqCst); let request = JsonRpcRequest::new(id, method, params); let (tx, rx) = oneshot::channel(); self.pending_requests.write().insert( id, PendingRequest { sender: tx, inline_callback, }, ); // RAII guard that removes the pending entry if this future is // dropped before the response arrives. Disarmed below before the // success return so the read loop owns the cleanup on the happy // path. let mut guard = PendingGuard { map: &self.pending_requests, id, armed: true, }; // The PendingGuard's drop removes the entry on every error path // and on cancellation; disarmed below before the success return so // the read loop owns the cleanup on the happy path. if let Err(error) = self.write(&request).await { warn!( elapsed_ms = request_start.elapsed().as_millis(), method = %method, request_id = id, status = "failed", error = %error, "JsonRpcClient::send_request JSON-RPC request finished" ); return Err(error); } let response = match rx.await { Ok(response) => response, Err(_) => { let error = ErrorKind::Protocol(ProtocolErrorKind::RequestCancelled).into(); warn!( elapsed_ms = request_start.elapsed().as_millis(), method = %method, request_id = id, status = "failed", error = %error, "JsonRpcClient::send_request JSON-RPC request finished" ); return Err(error); } }; guard.disarm(); if let Some(error) = &response.error { warn!( elapsed_ms = request_start.elapsed().as_millis(), method = %method, request_id = id, status = "failed", code = error.code, error = %error.message, "JsonRpcClient::send_request JSON-RPC request finished" ); } else { debug!( elapsed_ms = request_start.elapsed().as_millis(), method = %method, request_id = id, status = "succeeded", "JsonRpcClient::send_request JSON-RPC request finished" ); } Ok(response) } /// Write a Content-Length-framed JSON-RPC message to the transport. /// /// # Cancel safety /// /// **Cancel-safe.** Pre-serializes the body, enqueues it on the writer /// actor's command channel, and awaits an ack. Caller cancellation /// drops the ack receiver; the actor still completes the frame and /// flushes. A partial frame can never appear on the wire. pub async fn write(&self, message: &T) -> Result<(), Error> { let body = serde_json::to_vec(message)?; let mut frame = Vec::with_capacity(CONTENT_LENGTH_HEADER.len() + 16 + body.len() + 4); frame.extend_from_slice(CONTENT_LENGTH_HEADER.as_bytes()); frame.extend_from_slice(body.len().to_string().as_bytes()); frame.extend_from_slice(b"\r\n\r\n"); frame.extend_from_slice(&body); let (ack_tx, ack_rx) = oneshot::channel(); self.write_tx .send(WriteCommand { frame, ack: ack_tx }) .map_err(|_| { Error::from(std::io::Error::new( std::io::ErrorKind::BrokenPipe, "writer actor has shut down", )) })?; match ack_rx.await { Ok(Ok(())) => Ok(()), Ok(Err(e)) => Err(Error::from(e)), Err(_) => Err(Error::from(std::io::Error::new( std::io::ErrorKind::BrokenPipe, "writer actor dropped ack without responding", ))), } } } /// RAII guard that removes a pending-request entry from the map if the /// owning future is dropped before the response arrives. Disarmed on the /// happy path so the read loop's response handling owns the cleanup. struct PendingGuard<'a> { map: &'a RwLock>, id: u64, armed: bool, } impl PendingGuard<'_> { fn disarm(&mut self) { self.armed = false; } } impl Drop for PendingGuard<'_> { fn drop(&mut self) { if self.armed { self.map.write().remove(&self.id); } } } #[cfg(test)] mod tests { use super::*; #[test] fn deserialize_notification() { let json = r#"{"jsonrpc":"2.0","method":"session.event","params":{"id":"e1"}}"#; let msg: JsonRpcMessage = serde_json::from_str(json).unwrap(); assert!(matches!(msg, JsonRpcMessage::Notification(n) if n.method == "session.event")); } #[test] fn deserialize_request() { let json = r#"{"jsonrpc":"2.0","id":5,"method":"permission.request","params":{"kind":"shell"}}"#; let msg: JsonRpcMessage = serde_json::from_str(json).unwrap(); assert!( matches!(msg, JsonRpcMessage::Request(r) if r.id == 5 && r.method == "permission.request") ); } #[test] fn deserialize_response_with_result() { let json = r#"{"jsonrpc":"2.0","id":3,"result":{"ok":true}}"#; let msg: JsonRpcMessage = serde_json::from_str(json).unwrap(); assert!(matches!(msg, JsonRpcMessage::Response(r) if r.id == 3 && !r.is_error())); } #[test] fn deserialize_error_response() { let json = r#"{"jsonrpc":"2.0","id":7,"error":{"code":-32600,"message":"Invalid Request"}}"#; let msg: JsonRpcMessage = serde_json::from_str(json).unwrap(); match msg { JsonRpcMessage::Response(r) => { assert!(r.is_error()); let err = r.error.unwrap(); assert_eq!(err.code, -32600); assert_eq!(err.message, "Invalid Request"); } other => panic!("expected Response, got {other:?}"), } } #[test] fn deserialize_rejects_non_object() { let result = serde_json::from_str::(r#""not an object""#); assert!(result.is_err()); } #[test] fn request_new_sets_version() { let req = JsonRpcRequest::new(42, "test.method", None); assert_eq!(req.jsonrpc, "2.0"); assert_eq!(req.id, 42); assert_eq!(req.method, "test.method"); assert!(req.params.is_none()); } #[test] fn request_serializes_camel_case() { let req = JsonRpcRequest::new(1, "ping", Some(serde_json::json!({}))); let json = serde_json::to_string(&req).unwrap(); assert!(json.contains(r#""jsonrpc":"2.0""#)); assert!(json.contains(r#""id":1"#)); assert!(json.contains(r#""method":"ping""#)); } #[test] fn notification_without_params_omits_field() { let n = JsonRpcNotification { jsonrpc: "2.0".into(), method: "ping".into(), params: None, }; let json = serde_json::to_string(&n).unwrap(); assert!(!json.contains("params")); } #[test] fn response_without_error_omits_field() { let r = JsonRpcResponse { jsonrpc: "2.0".into(), id: 1, result: Some(serde_json::json!(true)), error: None, }; let json = serde_json::to_string(&r).unwrap(); assert!(!json.contains("error")); } }