Skip to main content

github_copilot_sdk/
jsonrpc.rs

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
17/// Callback invoked synchronously by the JSON-RPC read loop the instant a
18/// successful response is parsed, before the response is delivered to the
19/// awaiter and before the read loop dispatches the next message. Use this
20/// when client-side state (for example, registering a server-assigned
21/// session id with the router) must be visible to any subsequent
22/// notification on the same connection.
23///
24/// If the callback returns an error, that error is delivered to the
25/// awaiter in place of the response.
26pub(crate) type InlineResponseCallback =
27    Box<dyn FnOnce(&JsonRpcResponse) -> Result<(), Error> + Send + Sync>;
28
29/// Internal pairing of the response delivery channel with an optional
30/// inline callback that the read loop runs synchronously before delivery.
31struct PendingRequest {
32    sender: oneshot::Sender<JsonRpcResponse>,
33    inline_callback: Option<InlineResponseCallback>,
34}
35
36/// A JSON-RPC 2.0 request message.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38#[serde(rename_all = "camelCase")]
39pub struct JsonRpcRequest {
40    /// Protocol version (always `"2.0"`).
41    pub jsonrpc: String,
42    /// Request ID for correlating responses.
43    pub id: u64,
44    /// RPC method name.
45    pub method: String,
46    /// Optional method parameters.
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub params: Option<Value>,
49}
50
51/// A JSON-RPC 2.0 response message.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct JsonRpcResponse {
55    /// Protocol version (always `"2.0"`).
56    pub jsonrpc: String,
57    /// Request ID this response correlates to.
58    pub id: u64,
59    /// Success payload (mutually exclusive with `error`).
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub result: Option<Value>,
62    /// Error payload (mutually exclusive with `result`).
63    #[serde(skip_serializing_if = "Option::is_none")]
64    pub error: Option<JsonRpcError>,
65}
66
67/// A JSON-RPC 2.0 error object.
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct JsonRpcError {
70    /// Numeric error code.
71    pub code: i32,
72    /// Human-readable error description.
73    pub message: String,
74    /// Optional structured error data.
75    #[serde(skip_serializing_if = "Option::is_none")]
76    pub data: Option<Value>,
77}
78
79/// Standard JSON-RPC 2.0 error codes.
80pub mod error_codes {
81    /// Method not found (-32601).
82    pub const METHOD_NOT_FOUND: i32 = -32601;
83    /// Invalid method parameters (-32602).
84    pub const INVALID_PARAMS: i32 = -32602;
85    /// Internal server error (-32603).
86    #[allow(dead_code, reason = "standard JSON-RPC code, reserved for future use")]
87    pub const INTERNAL_ERROR: i32 = -32603;
88}
89
90/// A JSON-RPC 2.0 notification (no `id`, no response expected).
91#[derive(Debug, Clone, Serialize, Deserialize)]
92#[serde(rename_all = "camelCase")]
93pub struct JsonRpcNotification {
94    /// Protocol version (always `"2.0"`).
95    pub jsonrpc: String,
96    /// Notification method name.
97    pub method: String,
98    /// Optional notification parameters.
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub params: Option<Value>,
101}
102
103/// A parsed JSON-RPC 2.0 message — request, response, or notification.
104#[derive(Debug, Clone, Serialize)]
105pub enum JsonRpcMessage {
106    /// An incoming or outgoing request.
107    Request(JsonRpcRequest),
108    /// A response to a previous request.
109    Response(JsonRpcResponse),
110    /// A fire-and-forget notification.
111    Notification(JsonRpcNotification),
112}
113
114/// Custom deserializer that dispatches based on field presence instead of
115/// `#[serde(untagged)]` which tries each variant sequentially (3× parse
116/// attempts for Notification — the hot-path streaming variant).
117///
118/// Dispatch logic:
119/// - has `id` + has `method` → Request
120/// - has `id` + no `method` → Response
121/// - no `id`                → Notification
122impl<'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        // Preserve the owned payload instead of rebuilding its JSON containers
136        // while serde validates the envelope. Optional null payloads remain None.
137        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    /// Create a new JSON-RPC request with the given ID, method, and params.
171    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    /// Returns `true` if this response contains an error.
183    #[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
191/// Rewrites unpaired UTF-16 surrogate escapes to `\uFFFD`.
192///
193/// Returns `None` when the body contains no unpaired surrogate, so valid
194/// frames do not incur a repair allocation.
195fn 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            // Consume non-Unicode escapes whole so an escaped backslash cannot
221            // be mistaken for the start of a surrogate escape.
222            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
253/// One framed JSON-RPC message handed to the writer actor.
254///
255/// `frame` is the fully serialized bytes (header + body); the caller pays
256/// the serde cost synchronously before enqueueing so the actor never sees a
257/// `Result` from JSON encoding. `ack` resolves once the bytes have been
258/// fully written and flushed (or the underlying I/O reports an error). If
259/// the caller drops the `oneshot::Receiver`, the actor still completes the
260/// frame — caller cancellation cannot desync the wire.
261struct WriteCommand {
262    frame: Vec<u8>,
263    ack: oneshot::Sender<Result<(), std::io::Error>>,
264}
265
266/// Low-level JSON-RPC 2.0 client over Content-Length-framed streams.
267///
268/// # Cancel safety
269///
270/// All public methods (`write`, `send_request`) are **cancel-safe**: the
271/// actual bytes hit the wire on a dedicated background actor task, so
272/// dropping the caller's future after `await` returns `Pending` cannot
273/// produce a partial frame on the wire. Frames either land atomically or
274/// the underlying I/O fails. See `cancel-safety review` artifact for the
275/// full RFD-400 reasoning.
276pub struct JsonRpcClient {
277    request_id: AtomicU64,
278    /// Sender side of the writer actor's command queue. Public methods
279    /// pre-serialize their frames and enqueue here; the background actor
280    /// drains the queue and serializes writes onto the underlying
281    /// `AsyncWrite`. Unbounded by design — RFD 400 explicitly permits this
282    /// for cancel-safety, and JSON-RPC frames are small relative to the
283    /// natural request/response back-pressure of the wire.
284    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    /// Create a new client from async read/write streams.
295    ///
296    /// Spawns two background tasks: a reader that dispatches incoming
297    /// messages to pending request channels, the notification broadcast,
298    /// or the request-forwarding channel; and a writer actor that owns the
299    /// underlying `AsyncWrite` and serializes frames atomically.
300    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    /// Writer-actor task. Owns the `AsyncWrite`, drains the command queue,
362    /// and writes each frame atomically (header + body + flush) before
363    /// signaling the ack.
364    ///
365    /// Caller-side cancellation cannot interrupt a write in progress:
366    /// dropping the ack `oneshot::Receiver` does not cancel the in-flight
367    /// I/O. Once `WriteCommand` is enqueued the frame is committed to land
368    /// on the wire (or surface an `io::Error` to the ack receiver if the
369    /// transport is broken).
370    ///
371    /// Exits cleanly when all senders drop (channel closes), flushing any
372    /// final buffered bytes.
373    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            // Caller may have dropped the ack receiver (e.g. their
386            // `await` was cancelled); that's fine — we still completed
387            // the write, which was the whole point.
388            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                            // Run the inline callback synchronously on the
412                            // read loop so any state it mutates (e.g.
413                            // registering a server-assigned session id with
414                            // the router) is visible before the loop reads
415                            // and dispatches the next message.
416                            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        // Drain in-flight requests so callers observe cancellation
477        // instead of hanging on a oneshot receiver.
478        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                // Dropping an undecodable frame could leave its pending
525                // request waiting forever because this layer has no timeout.
526                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    /// Send a JSON-RPC request and wait for the matching response.
544    ///
545    /// # Cancel safety
546    ///
547    /// **Cancel-safe.** The frame is committed to the wire via the writer
548    /// actor before this future yields; cancelling the await drops the
549    /// response oneshot but does not desync the transport. The pending-
550    /// requests map is cleaned up automatically (the `PendingGuard` drop
551    /// removes the entry, and the read loop's response handling tolerates
552    /// a missing entry).
553    #[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    /// Send a JSON-RPC request whose response is observed synchronously
564    /// by the read loop *before* it is delivered to the awaiter.
565    ///
566    /// The optional `inline_callback` runs on the JSON-RPC read task the
567    /// instant a successful response is parsed, and before the read loop
568    /// dispatches the next message. This is the only way to perform
569    /// client-side bookkeeping (for example, registering a server-
570    /// assigned session id with the router) that must be visible to any
571    /// notification or request that the server may emit on the same
572    /// connection immediately after the response.
573    ///
574    /// If the callback returns an error or panics, that error is
575    /// surfaced to the awaiter in place of the original response (the
576    /// response payload is discarded and an internal-error JSON-RPC
577    /// error is delivered instead). The error is never propagated back
578    /// to the server and does not crash the read loop.
579    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        // RAII guard that removes the pending entry if this future is
599        // dropped before the response arrives. Disarmed below before the
600        // success return so the read loop owns the cleanup on the happy
601        // path.
602        let mut guard = PendingGuard {
603            map: &self.pending_requests,
604            id,
605            armed: true,
606        };
607
608        // The PendingGuard's drop removes the entry on every error path
609        // and on cancellation; disarmed below before the success return so
610        // the read loop owns the cleanup on the happy path.
611        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    /// Write a Content-Length-framed JSON-RPC message to the transport.
662    ///
663    /// # Cancel safety
664    ///
665    /// **Cancel-safe.** Pre-serializes the body, enqueues it on the writer
666    /// actor's command channel, and awaits an ack. Caller cancellation
667    /// drops the ack receiver; the actor still completes the frame and
668    /// flushes. A partial frame can never appear on the wire.
669    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
698/// RAII guard that removes a pending-request entry from the map if the
699/// owning future is dropped before the response arrives. Disarmed on the
700/// happy path so the read loop's response handling owns the cleanup.
701struct 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}