agent_client_protocol/jsonrpc/
transport_actor.rs

1use std::pin::pin;
2
3// Types re-exported from crate root
4use crate::RawJsonRpcResponse as Response;
5use crate::jsonrpc::{RawJsonRpcMessage, TransportBatch, TransportBatchEntry, TransportFrame};
6use futures::StreamExt as _;
7use futures::channel::mpsc;
8use serde::Deserialize as _;
9
10enum ParsedIncomingLine {
11    Single(RawJsonRpcMessage),
12    Malformed { raw: String, error: crate::Error },
13    Batch(TransportBatch),
14}
15
16fn parse_incoming_line(line: &str) -> ParsedIncomingLine {
17    let value = match serde_json::from_str::<serde_json::Value>(line) {
18        Ok(value) => value,
19        Err(error) => {
20            // Serde reasons can quote input, so log only the category and position.
21            tracing::debug!(
22                category = ?error.classify(),
23                line = error.line(),
24                column = error.column(),
25                "Failed to parse incoming JSON-RPC JSON"
26            );
27            return ParsedIncomingLine::Malformed {
28                raw: line.to_owned(),
29                error: crate::Error::parse_error().data(serde_json::json!({ "line": line })),
30            };
31        }
32    };
33
34    match value {
35        serde_json::Value::Array(entries) if entries.is_empty() => ParsedIncomingLine::Malformed {
36            raw: line.to_owned(),
37            error: crate::Error::invalid_request(),
38        },
39        serde_json::Value::Array(entries) => {
40            let entries = entries
41                .into_iter()
42                .map(|entry| match RawJsonRpcMessage::deserialize(&entry) {
43                    Ok(message) => TransportBatchEntry::message(message),
44                    Err(error) => {
45                        tracing::debug!(
46                            category = ?error.classify(),
47                            "Invalid JSON-RPC batch entry"
48                        );
49                        TransportBatchEntry::malformed(entry, crate::Error::invalid_request())
50                    }
51                })
52                .collect::<Vec<_>>();
53
54            ParsedIncomingLine::Batch(
55                TransportBatch::from_entries(entries)
56                    .expect("a parsed non-empty JSON array retains at least one entry"),
57            )
58        }
59        value => match serde_json::from_value(value) {
60            Ok(message) => ParsedIncomingLine::Single(message),
61            Err(error) => {
62                tracing::debug!(category = ?error.classify(), "Invalid JSON-RPC message");
63                ParsedIncomingLine::Malformed {
64                    raw: line.to_owned(),
65                    error: crate::Error::invalid_request(),
66                }
67            }
68        },
69    }
70}
71
72impl TransportFrame {
73    /// Parse one JSON-RPC wire value while preserving batch boundaries.
74    ///
75    /// Every malformed value is retained as an explicit malformed frame or
76    /// batch entry. Standalone malformed input keeps its original text; batch
77    /// entries keep their parsed JSON values, source order, and batch boundary,
78    /// though reserialization may normalize whitespace. Protocol actors decide
79    /// whether a malformed value is call-shaped and requires an Error Response.
80    #[must_use]
81    pub fn parse_json(input: &str) -> Self {
82        match parse_incoming_line(input) {
83            ParsedIncomingLine::Single(message) => Self::Single(message),
84            ParsedIncomingLine::Malformed { raw, error } => Self::Malformed { raw, error },
85            ParsedIncomingLine::Batch(batch) => Self::Batch(batch),
86        }
87    }
88
89    /// Serialize this frame to its JSON-RPC wire representation.
90    ///
91    /// # Errors
92    ///
93    /// Returns an internal error if a valid message or batch cannot be
94    /// serialized. Malformed frames return their original wire text unchanged.
95    pub fn to_json(&self) -> Result<String, crate::Error> {
96        match self {
97            Self::Single(message) => {
98                serde_json::to_string(message).map_err(crate::Error::into_internal_error)
99            }
100            Self::Malformed { raw, .. } => Ok(raw.clone()),
101            Self::Batch(batch) => {
102                serde_json::to_string(batch).map_err(crate::Error::into_internal_error)
103            }
104        }
105    }
106}
107
108/// Transport outgoing actor for line streams: serializes [`TransportFrame`] values and yields
109/// lines.
110///
111/// This is a line-based variant of `transport_outgoing_actor` that works with a Sink<String>
112/// instead of an AsyncWrite byte stream. This enables interception of lines before they are
113/// written to the underlying transport.
114///
115/// This actor handles transport mechanics:
116/// - Serializes single messages, malformed values, and batches to JSON strings
117/// - Yields newline-terminated strings
118/// - Handles serialization errors
119///
120/// This is the transport layer - it has no knowledge of protocol semantics (IDs, correlation, etc.).
121async fn transport_outgoing_frames_actor(
122    transport_rx: impl futures::Stream<Item = TransportFrame>,
123    outgoing_lines: impl futures::Sink<String, Error = std::io::Error>,
124) -> Result<(), crate::Error> {
125    use futures::SinkExt;
126    let mut transport_rx = pin!(transport_rx);
127    let mut outgoing_lines = pin!(outgoing_lines);
128
129    while let Some(frame) = transport_rx.next().await {
130        let json_rpc_message = match frame {
131            TransportFrame::Single(message) => message,
132            TransportFrame::Malformed { raw, .. } => {
133                let raw = malformed_line_value(raw)?;
134                tracing::trace!(
135                    message = %crate::util::redacted_json_line(&raw),
136                    "Relaying invalid JSON-RPC value"
137                );
138                outgoing_lines
139                    .send(raw)
140                    .await
141                    .map_err(crate::Error::into_internal_error)?;
142                continue;
143            }
144            TransportFrame::Batch(batch) => {
145                let line =
146                    serde_json::to_string(&batch).map_err(crate::Error::into_internal_error)?;
147                tracing::trace!(message = %crate::util::redacted_json_line(&line), "Sending JSON-RPC batch");
148                outgoing_lines
149                    .send(line)
150                    .await
151                    .map_err(crate::Error::into_internal_error)?;
152                continue;
153            }
154        };
155        match serde_json::to_string(&json_rpc_message) {
156            Ok(line) => {
157                tracing::trace!(message = %crate::util::redacted_json_line(&line), "Sending JSON-RPC message");
158                outgoing_lines
159                    .send(line)
160                    .await
161                    .map_err(crate::Error::into_internal_error)?;
162            }
163
164            Err(serialization_error) => {
165                match json_rpc_message {
166                    RawJsonRpcMessage::Request(_) | RawJsonRpcMessage::Notification(_) => {
167                        // If we failed to serialize a request,
168                        // just ignore it.
169                        //
170                        // Q: (Maybe it'd be nice to "reply" with an error?)
171                        tracing::error!(
172                            ?serialization_error,
173                            "Failed to serialize request, ignoring"
174                        );
175                    }
176                    RawJsonRpcMessage::Response(response) => {
177                        // If we failed to serialize a *response*,
178                        // send an error in response.
179                        let id = match response {
180                            Response::Result { id, .. } | Response::Error { id, .. } => id,
181                        };
182                        tracing::error!(
183                            ?serialization_error,
184                            ?id,
185                            "Failed to serialize response, sending internal_error instead"
186                        );
187                        let error_line = serde_json::to_string(&RawJsonRpcMessage::response(
188                            id,
189                            Err(crate::Error::internal_error()),
190                        ))
191                        .unwrap();
192                        outgoing_lines
193                            .send(error_line)
194                            .await
195                            .map_err(crate::Error::into_internal_error)?;
196                    }
197                }
198            }
199        }
200    }
201    outgoing_lines
202        .close()
203        .await
204        .map_err(crate::Error::into_internal_error)
205}
206
207/// A newline writer whose sink close reaches the actual write half. An unfold
208/// sink can flush each line, but has no way to forward `AsyncWrite::close`.
209pub(super) struct LineWriter<W> {
210    writer: std::pin::Pin<Box<W>>,
211    bytes: Vec<u8>,
212    written: usize,
213    closing: bool,
214}
215
216impl<W> LineWriter<W> {
217    pub(super) fn new(writer: W) -> Self {
218        Self {
219            writer: Box::pin(writer),
220            bytes: Vec::new(),
221            written: 0,
222            closing: false,
223        }
224    }
225}
226
227impl<W: futures::AsyncWrite> futures::Sink<String> for LineWriter<W> {
228    type Error = std::io::Error;
229
230    fn poll_ready(
231        self: std::pin::Pin<&mut Self>,
232        cx: &mut std::task::Context<'_>,
233    ) -> std::task::Poll<Result<(), Self::Error>> {
234        self.poll_flush(cx)
235    }
236
237    fn start_send(self: std::pin::Pin<&mut Self>, line: String) -> Result<(), Self::Error> {
238        let this = self.get_mut();
239        this.bytes = line.into_bytes();
240        this.bytes.push(b'\n');
241        this.written = 0;
242        Ok(())
243    }
244
245    fn poll_flush(
246        self: std::pin::Pin<&mut Self>,
247        cx: &mut std::task::Context<'_>,
248    ) -> std::task::Poll<Result<(), Self::Error>> {
249        let this = self.get_mut();
250        while this.written < this.bytes.len() {
251            let count = futures::ready!(
252                this.writer
253                    .as_mut()
254                    .poll_write(cx, &this.bytes[this.written..])
255            )?;
256            if count == 0 {
257                return std::task::Poll::Ready(Err(std::io::ErrorKind::WriteZero.into()));
258            }
259            this.written += count;
260        }
261        this.bytes.clear();
262        this.written = 0;
263        this.writer.as_mut().poll_flush(cx)
264    }
265
266    fn poll_close(
267        mut self: std::pin::Pin<&mut Self>,
268        cx: &mut std::task::Context<'_>,
269    ) -> std::task::Poll<Result<(), Self::Error>> {
270        if !self.closing {
271            futures::ready!(self.as_mut().poll_flush(cx))?;
272            self.closing = true;
273        }
274        self.get_mut().writer.as_mut().poll_close(cx)
275    }
276}
277
278fn malformed_line_value(raw: String) -> Result<String, crate::Error> {
279    if !raw.contains('\r') && !raw.contains('\n') {
280        return Ok(raw);
281    }
282
283    match serde_json::from_str::<serde_json::Value>(&raw) {
284        Ok(value) => serde_json::to_string(&value),
285        Err(_) => serde_json::to_string(&raw),
286    }
287    .map_err(crate::Error::into_internal_error)
288}
289
290pub(super) async fn transport_outgoing_lines_actor(
291    transport_rx: impl futures::Stream<Item = TransportFrame>,
292    outgoing_lines: impl futures::Sink<String, Error = std::io::Error>,
293) -> Result<(), crate::Error> {
294    transport_outgoing_frames_actor(transport_rx, outgoing_lines).await
295}
296
297/// Transport incoming actor for line streams: parses lines into [`TransportFrame`] values.
298///
299/// This is a line-based variant of `transport_incoming_actor` that works with a
300/// Stream<Item = io::Result<String>> instead of an AsyncRead byte stream. This enables
301/// interception of lines before they are parsed.
302///
303/// This actor handles transport mechanics:
304/// - Reads lines from the stream
305/// - Parses individual messages and retains batch arrays in entry order
306/// - Handles malformed JSON, empty batches, and invalid batch entries
307///
308/// This is the transport layer - it has no knowledge of protocol semantics.
309pub(super) async fn transport_incoming_lines_actor(
310    incoming_lines: impl futures::Stream<Item = std::io::Result<String>>,
311    transport_tx: mpsc::UnboundedSender<TransportFrame>,
312) -> Result<(), crate::Error> {
313    let mut incoming_lines = pin!(incoming_lines);
314    while let Some(line_result) = incoming_lines.next().await {
315        let line = line_result.map_err(crate::Error::into_internal_error)?;
316        tracing::trace!(message = %crate::util::redacted_json_line(&line), "Received JSON-RPC message");
317
318        match parse_incoming_line(&line) {
319            ParsedIncomingLine::Single(message) => {
320                transport_tx
321                    .unbounded_send(TransportFrame::Single(message))
322                    .map_err(crate::Error::into_internal_error)?;
323            }
324            ParsedIncomingLine::Malformed { raw, error } => {
325                transport_tx
326                    .unbounded_send(TransportFrame::Malformed { raw, error })
327                    .map_err(crate::Error::into_internal_error)?;
328            }
329            ParsedIncomingLine::Batch(entries) => {
330                transport_tx
331                    .unbounded_send(TransportFrame::Batch(entries))
332                    .map_err(crate::Error::into_internal_error)?;
333            }
334        }
335    }
336    Ok(())
337}
338
339#[cfg(test)]
340mod tests {
341    use std::sync::{Arc, Mutex};
342
343    use super::*;
344    use crate::ErrorCode;
345
346    #[derive(Default)]
347    struct PendingCloseWriter {
348        bytes: Vec<u8>,
349        close_polls: usize,
350    }
351
352    impl futures::AsyncWrite for PendingCloseWriter {
353        fn poll_write(
354            mut self: std::pin::Pin<&mut Self>,
355            _: &mut std::task::Context<'_>,
356            bytes: &[u8],
357        ) -> std::task::Poll<std::io::Result<usize>> {
358            self.bytes.extend_from_slice(bytes);
359            std::task::Poll::Ready(Ok(bytes.len()))
360        }
361
362        fn poll_flush(
363            self: std::pin::Pin<&mut Self>,
364            _: &mut std::task::Context<'_>,
365        ) -> std::task::Poll<std::io::Result<()>> {
366            assert_eq!(self.close_polls, 0, "do not flush again during shutdown");
367            std::task::Poll::Ready(Ok(()))
368        }
369
370        fn poll_close(
371            mut self: std::pin::Pin<&mut Self>,
372            cx: &mut std::task::Context<'_>,
373        ) -> std::task::Poll<std::io::Result<()>> {
374            self.close_polls += 1;
375            if self.close_polls == 1 {
376                cx.waker().wake_by_ref();
377                std::task::Poll::Pending
378            } else {
379                std::task::Poll::Ready(Ok(()))
380            }
381        }
382    }
383
384    #[test]
385    fn byte_writer_flushes_then_continues_pending_shutdown_without_reflushing() {
386        use futures::SinkExt as _;
387        let mut sink = LineWriter::new(PendingCloseWriter::default());
388        futures::executor::block_on(async {
389            sink.send("line".to_string()).await.unwrap();
390            sink.close().await.unwrap();
391        });
392        assert_eq!(sink.writer.bytes, b"line\n");
393        assert_eq!(sink.writer.close_polls, 2);
394    }
395
396    #[test]
397    fn parses_batch_entries_independently() {
398        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
399            r#"[
400                {"jsonrpc":"2.0","id":1,"method":"one","params":{}},
401                17,
402                {"jsonrpc":"2.0","method":"two","params":{}}
403            ]"#,
404        ) else {
405            panic!("expected a JSON-RPC batch");
406        };
407
408        let entries = batch.iter_results().collect::<Vec<_>>();
409        assert_eq!(entries.len(), 3);
410        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Request(_))));
411        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
412        assert!(matches!(entries[2], Ok(RawJsonRpcMessage::Notification(_))));
413    }
414
415    #[test]
416    fn preserves_every_invalid_member_of_response_batches() {
417        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
418            r#"[
419                {"jsonrpc":"2.0","id":1,"result":{"ok":true}},
420                17,
421                {"jsonrpc":"2.0","id":2,"result":null,"error":{"code":-32603,"message":"Internal error"}},
422                {"jsonrpc":"2.0","id":3,"error":{"code":-32603,"message":"Internal error"}}
423            ]"#,
424        ) else {
425            panic!("expected a JSON-RPC batch");
426        };
427
428        let entries = batch.iter_results().collect::<Vec<_>>();
429        assert_eq!(entries.len(), 4);
430        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Response(_))));
431        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
432        assert_eq!(entries[2].unwrap_err().code, ErrorCode::InvalidRequest);
433        assert!(matches!(entries[3], Ok(RawJsonRpcMessage::Response(_))));
434    }
435
436    #[test]
437    fn preserves_invalid_value_beside_malformed_response() {
438        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
439            r#"[
440                17,
441                {"jsonrpc":"2.0","id":1,"result":null,"error":{"code":-32603,"message":"Internal error"}}
442            ]"#,
443        ) else {
444            panic!("expected a JSON-RPC batch");
445        };
446
447        let entries = batch.iter_results().collect::<Vec<_>>();
448        assert_eq!(entries.len(), 2);
449        assert_eq!(entries[0].unwrap_err().code, ErrorCode::InvalidRequest);
450        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
451    }
452
453    #[test]
454    fn preserves_entirely_malformed_response_shaped_batch() {
455        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
456            r#"[
457                {"jsonrpc":"2.0","id":1,"result":null,"error":{"code":-32603,"message":"Internal error"}}
458            ]"#,
459        ) else {
460            panic!("expected a retained JSON-RPC batch");
461        };
462
463        let entries = batch.iter_results().collect::<Vec<_>>();
464        assert_eq!(entries.len(), 1);
465        assert_eq!(entries[0].unwrap_err().code, ErrorCode::InvalidRequest);
466    }
467
468    #[test]
469    fn preserves_invalid_call_shaped_member_beside_response() {
470        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
471            r#"[
472                {"jsonrpc":"2.0","id":1,"result":null},
473                {"jsonrpc":"2.0","method":1}
474            ]"#,
475        ) else {
476            panic!("expected a JSON-RPC batch");
477        };
478
479        let entries = batch.iter_results().collect::<Vec<_>>();
480        assert_eq!(entries.len(), 2);
481        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Response(_))));
482        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
483    }
484
485    #[test]
486    fn preserves_malformed_response_shaped_member_beside_request() {
487        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
488            r#"[
489                {"jsonrpc":"2.0","id":1,"method":"one","params":{}},
490                {"jsonrpc":"2.0","id":2,"result":null,"error":{"code":-32603,"message":"Internal error"}}
491            ]"#,
492        ) else {
493            panic!("expected a JSON-RPC batch");
494        };
495
496        let entries = batch.iter_results().collect::<Vec<_>>();
497        assert_eq!(entries.len(), 2);
498        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Request(_))));
499        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
500    }
501
502    #[test]
503    fn preserves_malformed_standalone_response() {
504        let ParsedIncomingLine::Malformed { error, .. } = parse_incoming_line(
505            r#"{"jsonrpc":"2.0","id":1,"result":null,"error":{"code":-32603,"message":"Internal error"}}"#,
506        ) else {
507            panic!("expected one retained invalid response");
508        };
509
510        assert_eq!(error.code, ErrorCode::InvalidRequest);
511    }
512
513    #[test]
514    fn preserves_malformed_call_shaped_standalone_message() {
515        let ParsedIncomingLine::Malformed { error, .. } =
516            parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"method":"one","result":null}"#)
517        else {
518            panic!("expected one invalid-request error");
519        };
520
521        assert_eq!(error.code, ErrorCode::InvalidRequest);
522    }
523
524    #[test]
525    fn parses_valid_standalone_response() {
526        assert!(matches!(
527            parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"result":{"ok":true}}"#),
528            ParsedIncomingLine::Single(RawJsonRpcMessage::Response(_))
529        ));
530    }
531
532    #[test]
533    fn all_invalid_batch_defaults_to_call_errors() {
534        let ParsedIncomingLine::Batch(batch) = parse_incoming_line("[1, 2, 3]") else {
535            panic!("expected a JSON-RPC batch");
536        };
537
538        assert_eq!(batch.len(), 3);
539        assert!(
540            batch
541                .iter_results()
542                .all(|entry| entry.unwrap_err().code == ErrorCode::InvalidRequest)
543        );
544    }
545
546    #[test]
547    fn empty_batch_is_an_invalid_request() {
548        let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("[]") else {
549            panic!("expected one invalid-request error");
550        };
551
552        assert_eq!(raw, "[]");
553        assert_eq!(error.code, ErrorCode::InvalidRequest);
554    }
555
556    #[test]
557    fn malformed_json_is_a_parse_error() {
558        let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("[") else {
559            panic!("expected one parse error");
560        };
561
562        assert_eq!(raw, "[");
563        assert_eq!(error.code, ErrorCode::ParseError);
564    }
565
566    #[test]
567    fn valid_json_with_an_invalid_envelope_is_an_invalid_request() {
568        let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("17") else {
569            panic!("expected one invalid-request error");
570        };
571
572        assert_eq!(raw, "17");
573        assert_eq!(error.code, ErrorCode::InvalidRequest);
574    }
575
576    #[tokio::test]
577    async fn multiline_malformed_frame_is_written_as_one_line_value() {
578        let raw = "not json\r\n{\"jsonrpc\":\"2.0\",\"method\":\"injected\"}".to_string();
579        let captured = Arc::new(Mutex::new(Vec::new()));
580        let outgoing = futures::sink::unfold(captured.clone(), |captured, line| async move {
581            captured.lock().unwrap().push(line);
582            Ok::<_, std::io::Error>(captured)
583        });
584
585        transport_outgoing_frames_actor(
586            futures::stream::iter([TransportFrame::Malformed {
587                raw: raw.clone(),
588                error: crate::Error::parse_error(),
589            }]),
590            outgoing,
591        )
592        .await
593        .unwrap();
594
595        let lines = captured.lock().unwrap();
596        assert_eq!(lines.len(), 1);
597        assert!(!lines[0].contains('\r') && !lines[0].contains('\n'));
598        assert_eq!(serde_json::from_str::<String>(&lines[0]).unwrap(), raw);
599    }
600
601    #[test]
602    fn multiline_invalid_json_rpc_value_is_compacted_without_changing_value() {
603        let raw = "{\n  \"jsonrpc\": \"2.0\",\n  \"method\": 1\n}".to_string();
604        let expected = serde_json::from_str::<serde_json::Value>(&raw).unwrap();
605        let line = malformed_line_value(raw).unwrap();
606
607        assert!(!line.contains('\r') && !line.contains('\n'));
608        assert_eq!(
609            serde_json::from_str::<serde_json::Value>(&line).unwrap(),
610            expected
611        );
612    }
613}
agent_client_protocol/jsonrpc/transport_actor.rs

agent_client_protocol/jsonrpc/
transport_actor.rs
