Skip to main content

agent_client_protocol/jsonrpc/
transport_actor.rs

1use std::pin::pin;
2
3// Types re-exported from crate root
4use crate::jsonrpc::{RawJsonRpcMessage, TransportBatch, TransportBatchEntry, TransportFrame};
5use crate::schema::v1::Response;
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            tracing::debug!(?error, "Failed to parse incoming JSON-RPC JSON");
21            return ParsedIncomingLine::Malformed {
22                raw: line.to_owned(),
23                error: crate::Error::parse_error().data(serde_json::json!({ "line": line })),
24            };
25        }
26    };
27
28    match value {
29        serde_json::Value::Array(entries) if entries.is_empty() => ParsedIncomingLine::Malformed {
30            raw: line.to_owned(),
31            error: crate::Error::invalid_request(),
32        },
33        serde_json::Value::Array(entries) => {
34            let entries = entries
35                .into_iter()
36                .map(|entry| match RawJsonRpcMessage::deserialize(&entry) {
37                    Ok(message) => TransportBatchEntry::message(message),
38                    Err(error) => {
39                        tracing::debug!(?error, "Invalid JSON-RPC batch entry");
40                        TransportBatchEntry::malformed(entry, crate::Error::invalid_request())
41                    }
42                })
43                .collect::<Vec<_>>();
44
45            ParsedIncomingLine::Batch(
46                TransportBatch::from_entries(entries)
47                    .expect("a parsed non-empty JSON array retains at least one entry"),
48            )
49        }
50        value => match serde_json::from_value(value) {
51            Ok(message) => ParsedIncomingLine::Single(message),
52            Err(error) => {
53                tracing::debug!(?error, "Invalid JSON-RPC message");
54                ParsedIncomingLine::Malformed {
55                    raw: line.to_owned(),
56                    error: crate::Error::invalid_request(),
57                }
58            }
59        },
60    }
61}
62
63impl TransportFrame {
64    /// Parse one JSON-RPC wire value while preserving batch boundaries.
65    ///
66    /// Every malformed value is retained as an explicit malformed frame or
67    /// batch entry. Standalone malformed input keeps its original text; batch
68    /// entries keep their parsed JSON values, source order, and batch boundary,
69    /// though reserialization may normalize whitespace. Protocol actors decide
70    /// whether a malformed value is call-shaped and requires an Error Response.
71    #[must_use]
72    pub fn parse_json(input: &str) -> Self {
73        match parse_incoming_line(input) {
74            ParsedIncomingLine::Single(message) => Self::Single(message),
75            ParsedIncomingLine::Malformed { raw, error } => Self::Malformed { raw, error },
76            ParsedIncomingLine::Batch(batch) => Self::Batch(batch),
77        }
78    }
79
80    /// Serialize this frame to its JSON-RPC wire representation.
81    ///
82    /// # Errors
83    ///
84    /// Returns an internal error if a valid message or batch cannot be
85    /// serialized. Malformed frames return their original wire text unchanged.
86    pub fn to_json(&self) -> Result<String, crate::Error> {
87        match self {
88            Self::Single(message) => {
89                serde_json::to_string(message).map_err(crate::Error::into_internal_error)
90            }
91            Self::Malformed { raw, .. } => Ok(raw.clone()),
92            Self::Batch(batch) => {
93                serde_json::to_string(batch).map_err(crate::Error::into_internal_error)
94            }
95        }
96    }
97}
98
99/// Transport outgoing actor for line streams: serializes [`TransportFrame`] values and yields
100/// lines.
101///
102/// This is a line-based variant of `transport_outgoing_actor` that works with a Sink<String>
103/// instead of an AsyncWrite byte stream. This enables interception of lines before they are
104/// written to the underlying transport.
105///
106/// This actor handles transport mechanics:
107/// - Serializes single messages, malformed values, and batches to JSON strings
108/// - Yields newline-terminated strings
109/// - Handles serialization errors
110///
111/// This is the transport layer - it has no knowledge of protocol semantics (IDs, correlation, etc.).
112async fn transport_outgoing_frames_actor(
113    transport_rx: impl futures::Stream<Item = TransportFrame>,
114    outgoing_lines: impl futures::Sink<String, Error = std::io::Error>,
115) -> Result<(), crate::Error> {
116    use futures::SinkExt;
117    let mut transport_rx = pin!(transport_rx);
118    let mut outgoing_lines = pin!(outgoing_lines);
119
120    while let Some(frame) = transport_rx.next().await {
121        let json_rpc_message = match frame {
122            TransportFrame::Single(message) => message,
123            TransportFrame::Malformed { raw, .. } => {
124                let raw = malformed_line_value(raw)?;
125                tracing::trace!(message = ?raw, "Relaying invalid JSON-RPC value");
126                outgoing_lines
127                    .send(raw)
128                    .await
129                    .map_err(crate::Error::into_internal_error)?;
130                continue;
131            }
132            TransportFrame::Batch(batch) => {
133                let line =
134                    serde_json::to_string(&batch).map_err(crate::Error::into_internal_error)?;
135                tracing::trace!(message = %line, "Sending JSON-RPC batch");
136                outgoing_lines
137                    .send(line)
138                    .await
139                    .map_err(crate::Error::into_internal_error)?;
140                continue;
141            }
142        };
143        match serde_json::to_string(&json_rpc_message) {
144            Ok(line) => {
145                tracing::trace!(message = %line, "Sending JSON-RPC message");
146                outgoing_lines
147                    .send(line)
148                    .await
149                    .map_err(crate::Error::into_internal_error)?;
150            }
151
152            Err(serialization_error) => {
153                match json_rpc_message {
154                    RawJsonRpcMessage::Request(_) | RawJsonRpcMessage::Notification(_) => {
155                        // If we failed to serialize a request,
156                        // just ignore it.
157                        //
158                        // Q: (Maybe it'd be nice to "reply" with an error?)
159                        tracing::error!(
160                            ?serialization_error,
161                            "Failed to serialize request, ignoring"
162                        );
163                    }
164                    RawJsonRpcMessage::Response(response) => {
165                        // If we failed to serialize a *response*,
166                        // send an error in response.
167                        let id = match response {
168                            Response::Result { id, .. } | Response::Error { id, .. } => id,
169                        };
170                        tracing::error!(
171                            ?serialization_error,
172                            ?id,
173                            "Failed to serialize response, sending internal_error instead"
174                        );
175                        let error_line = serde_json::to_string(&RawJsonRpcMessage::response(
176                            id,
177                            Err(crate::Error::internal_error()),
178                        ))
179                        .unwrap();
180                        outgoing_lines
181                            .send(error_line)
182                            .await
183                            .map_err(crate::Error::into_internal_error)?;
184                    }
185                }
186            }
187        }
188    }
189    Ok(())
190}
191
192fn malformed_line_value(raw: String) -> Result<String, crate::Error> {
193    if !raw.contains('\r') && !raw.contains('\n') {
194        return Ok(raw);
195    }
196
197    match serde_json::from_str::<serde_json::Value>(&raw) {
198        Ok(value) => serde_json::to_string(&value),
199        Err(_) => serde_json::to_string(&raw),
200    }
201    .map_err(crate::Error::into_internal_error)
202}
203
204pub(super) async fn transport_outgoing_lines_actor(
205    transport_rx: mpsc::UnboundedReceiver<TransportFrame>,
206    outgoing_lines: impl futures::Sink<String, Error = std::io::Error>,
207) -> Result<(), crate::Error> {
208    transport_outgoing_frames_actor(transport_rx, outgoing_lines).await
209}
210
211/// Transport incoming actor for line streams: parses lines into [`TransportFrame`] values.
212///
213/// This is a line-based variant of `transport_incoming_actor` that works with a
214/// Stream<Item = io::Result<String>> instead of an AsyncRead byte stream. This enables
215/// interception of lines before they are parsed.
216///
217/// This actor handles transport mechanics:
218/// - Reads lines from the stream
219/// - Parses individual messages and retains batch arrays in entry order
220/// - Handles malformed JSON, empty batches, and invalid batch entries
221///
222/// This is the transport layer - it has no knowledge of protocol semantics.
223pub(super) async fn transport_incoming_lines_actor(
224    incoming_lines: impl futures::Stream<Item = std::io::Result<String>>,
225    transport_tx: mpsc::UnboundedSender<TransportFrame>,
226) -> Result<(), crate::Error> {
227    let mut incoming_lines = pin!(incoming_lines);
228    while let Some(line_result) = incoming_lines.next().await {
229        let line = line_result.map_err(crate::Error::into_internal_error)?;
230        tracing::trace!(message = %line, "Received JSON-RPC message");
231
232        match parse_incoming_line(&line) {
233            ParsedIncomingLine::Single(message) => {
234                transport_tx
235                    .unbounded_send(TransportFrame::Single(message))
236                    .map_err(crate::Error::into_internal_error)?;
237            }
238            ParsedIncomingLine::Malformed { raw, error } => {
239                transport_tx
240                    .unbounded_send(TransportFrame::Malformed { raw, error })
241                    .map_err(crate::Error::into_internal_error)?;
242            }
243            ParsedIncomingLine::Batch(entries) => {
244                transport_tx
245                    .unbounded_send(TransportFrame::Batch(entries))
246                    .map_err(crate::Error::into_internal_error)?;
247            }
248        }
249    }
250    Ok(())
251}
252
253#[cfg(test)]
254mod tests {
255    use std::sync::{Arc, Mutex};
256
257    use super::*;
258    use crate::ErrorCode;
259
260    #[test]
261    fn parses_batch_entries_independently() {
262        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
263            r#"[
264                {"jsonrpc":"2.0","id":1,"method":"one","params":{}},
265                17,
266                {"jsonrpc":"2.0","method":"two","params":{}}
267            ]"#,
268        ) else {
269            panic!("expected a JSON-RPC batch");
270        };
271
272        let entries = batch.iter_results().collect::<Vec<_>>();
273        assert_eq!(entries.len(), 3);
274        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Request(_))));
275        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
276        assert!(matches!(entries[2], Ok(RawJsonRpcMessage::Notification(_))));
277    }
278
279    #[test]
280    fn preserves_every_invalid_member_of_response_batches() {
281        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
282            r#"[
283                {"jsonrpc":"2.0","id":1,"result":{"ok":true}},
284                17,
285                {"jsonrpc":"2.0","id":2,"result":null,"error":{"code":-32603,"message":"Internal error"}},
286                {"jsonrpc":"2.0","id":3,"error":{"code":-32603,"message":"Internal error"}}
287            ]"#,
288        ) else {
289            panic!("expected a JSON-RPC batch");
290        };
291
292        let entries = batch.iter_results().collect::<Vec<_>>();
293        assert_eq!(entries.len(), 4);
294        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Response(_))));
295        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
296        assert_eq!(entries[2].unwrap_err().code, ErrorCode::InvalidRequest);
297        assert!(matches!(entries[3], Ok(RawJsonRpcMessage::Response(_))));
298    }
299
300    #[test]
301    fn preserves_invalid_value_beside_malformed_response() {
302        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
303            r#"[
304                17,
305                {"jsonrpc":"2.0","id":1,"result":null,"error":{"code":-32603,"message":"Internal error"}}
306            ]"#,
307        ) else {
308            panic!("expected a JSON-RPC batch");
309        };
310
311        let entries = batch.iter_results().collect::<Vec<_>>();
312        assert_eq!(entries.len(), 2);
313        assert_eq!(entries[0].unwrap_err().code, ErrorCode::InvalidRequest);
314        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
315    }
316
317    #[test]
318    fn preserves_entirely_malformed_response_shaped_batch() {
319        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
320            r#"[
321                {"jsonrpc":"2.0","id":1,"result":null,"error":{"code":-32603,"message":"Internal error"}}
322            ]"#,
323        ) else {
324            panic!("expected a retained JSON-RPC batch");
325        };
326
327        let entries = batch.iter_results().collect::<Vec<_>>();
328        assert_eq!(entries.len(), 1);
329        assert_eq!(entries[0].unwrap_err().code, ErrorCode::InvalidRequest);
330    }
331
332    #[test]
333    fn preserves_invalid_call_shaped_member_beside_response() {
334        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
335            r#"[
336                {"jsonrpc":"2.0","id":1,"result":null},
337                {"jsonrpc":"2.0","method":1}
338            ]"#,
339        ) else {
340            panic!("expected a JSON-RPC batch");
341        };
342
343        let entries = batch.iter_results().collect::<Vec<_>>();
344        assert_eq!(entries.len(), 2);
345        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Response(_))));
346        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
347    }
348
349    #[test]
350    fn preserves_malformed_response_shaped_member_beside_request() {
351        let ParsedIncomingLine::Batch(batch) = parse_incoming_line(
352            r#"[
353                {"jsonrpc":"2.0","id":1,"method":"one","params":{}},
354                {"jsonrpc":"2.0","id":2,"result":null,"error":{"code":-32603,"message":"Internal error"}}
355            ]"#,
356        ) else {
357            panic!("expected a JSON-RPC batch");
358        };
359
360        let entries = batch.iter_results().collect::<Vec<_>>();
361        assert_eq!(entries.len(), 2);
362        assert!(matches!(entries[0], Ok(RawJsonRpcMessage::Request(_))));
363        assert_eq!(entries[1].unwrap_err().code, ErrorCode::InvalidRequest);
364    }
365
366    #[test]
367    fn preserves_malformed_standalone_response() {
368        let ParsedIncomingLine::Malformed { error, .. } = parse_incoming_line(
369            r#"{"jsonrpc":"2.0","id":1,"result":null,"error":{"code":-32603,"message":"Internal error"}}"#,
370        ) else {
371            panic!("expected one retained invalid response");
372        };
373
374        assert_eq!(error.code, ErrorCode::InvalidRequest);
375    }
376
377    #[test]
378    fn preserves_malformed_call_shaped_standalone_message() {
379        let ParsedIncomingLine::Malformed { error, .. } =
380            parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"method":"one","result":null}"#)
381        else {
382            panic!("expected one invalid-request error");
383        };
384
385        assert_eq!(error.code, ErrorCode::InvalidRequest);
386    }
387
388    #[test]
389    fn parses_valid_standalone_response() {
390        assert!(matches!(
391            parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"result":{"ok":true}}"#),
392            ParsedIncomingLine::Single(RawJsonRpcMessage::Response(_))
393        ));
394    }
395
396    #[test]
397    fn all_invalid_batch_defaults_to_call_errors() {
398        let ParsedIncomingLine::Batch(batch) = parse_incoming_line("[1, 2, 3]") else {
399            panic!("expected a JSON-RPC batch");
400        };
401
402        assert_eq!(batch.len(), 3);
403        assert!(
404            batch
405                .iter_results()
406                .all(|entry| entry.unwrap_err().code == ErrorCode::InvalidRequest)
407        );
408    }
409
410    #[test]
411    fn empty_batch_is_an_invalid_request() {
412        let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("[]") else {
413            panic!("expected one invalid-request error");
414        };
415
416        assert_eq!(raw, "[]");
417        assert_eq!(error.code, ErrorCode::InvalidRequest);
418    }
419
420    #[test]
421    fn malformed_json_is_a_parse_error() {
422        let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("[") else {
423            panic!("expected one parse error");
424        };
425
426        assert_eq!(raw, "[");
427        assert_eq!(error.code, ErrorCode::ParseError);
428    }
429
430    #[test]
431    fn valid_json_with_an_invalid_envelope_is_an_invalid_request() {
432        let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("17") else {
433            panic!("expected one invalid-request error");
434        };
435
436        assert_eq!(raw, "17");
437        assert_eq!(error.code, ErrorCode::InvalidRequest);
438    }
439
440    #[tokio::test]
441    async fn multiline_malformed_frame_is_written_as_one_line_value() {
442        let raw = "not json\r\n{\"jsonrpc\":\"2.0\",\"method\":\"injected\"}".to_string();
443        let captured = Arc::new(Mutex::new(Vec::new()));
444        let outgoing = futures::sink::unfold(captured.clone(), |captured, line| async move {
445            captured.lock().unwrap().push(line);
446            Ok::<_, std::io::Error>(captured)
447        });
448
449        transport_outgoing_frames_actor(
450            futures::stream::iter([TransportFrame::Malformed {
451                raw: raw.clone(),
452                error: crate::Error::parse_error(),
453            }]),
454            outgoing,
455        )
456        .await
457        .unwrap();
458
459        let lines = captured.lock().unwrap();
460        assert_eq!(lines.len(), 1);
461        assert!(!lines[0].contains('\r') && !lines[0].contains('\n'));
462        assert_eq!(serde_json::from_str::<String>(&lines[0]).unwrap(), raw);
463    }
464
465    #[test]
466    fn multiline_invalid_json_rpc_value_is_compacted_without_changing_value() {
467        let raw = "{\n  \"jsonrpc\": \"2.0\",\n  \"method\": 1\n}".to_string();
468        let expected = serde_json::from_str::<serde_json::Value>(&raw).unwrap();
469        let line = malformed_line_value(raw).unwrap();
470
471        assert!(!line.contains('\r') && !line.contains('\n'));
472        assert_eq!(
473            serde_json::from_str::<serde_json::Value>(&line).unwrap(),
474            expected
475        );
476    }
477}