agent_client_protocol/jsonrpc/

transport_actor.rs

use std::pin::pin;

// Types re-exported from crate root
use crate::RawJsonRpcResponse as Response;
use crate::jsonrpc::{RawJsonRpcMessage, TransportFrame};
use futures::StreamExt as _;
use futures::channel::mpsc;

fn parse_incoming_line(line: &str) -> TransportFrame {
    let value = match serde_json::from_str::<serde_json::Value>(line) {
        Ok(value) => value,
        Err(error) => {
            // Serde reasons can quote input, so log only the category and position.
            tracing::debug!(
                category = ?error.classify(),
                line = error.line(),
                column = error.column(),
                "Failed to parse incoming JSON-RPC JSON"
            );
            return TransportFrame::Malformed {
                raw: line.to_owned(),
                error: crate::Error::parse_error().data(serde_json::json!({ "line": line })),
            };
        }
    };

    match value {
        // JSON-RPC batching is not supported. An array, including an empty
        // one, is one invalid request answered with a single error.
        serde_json::Value::Array(_) => TransportFrame::Malformed {
            raw: line.to_owned(),
            error: crate::util::batch_not_supported(),
        },
        value => match serde_json::from_value(value) {
            Ok(message) => TransportFrame::Single(message),
            Err(error) => {
                tracing::debug!(category = ?error.classify(), "Invalid JSON-RPC message");
                TransportFrame::Malformed {
                    raw: line.to_owned(),
                    error: crate::Error::invalid_request(),
                }
            }
        },
    }
}

impl TransportFrame {
    /// Parse one JSON-RPC wire value.
    ///
    /// Every malformed value is retained as an explicit malformed frame with
    /// its original text. JSON-RPC batches are not supported: a JSON array is
    /// a malformed frame with an invalid-request (`-32600`) error. Protocol
    /// actors decide whether a malformed value is call-shaped and requires an
    /// Error Response.
    #[must_use]
    pub fn parse_json(input: &str) -> Self {
        parse_incoming_line(input)
    }

    /// Serialize this frame to its JSON-RPC wire representation.
    ///
    /// # Errors
    ///
    /// Returns an internal error if a valid message cannot be
    /// serialized. Malformed frames return their original wire text unchanged.
    pub fn to_json(&self) -> Result<String, crate::Error> {
        match self {
            Self::Single(message) => {
                serde_json::to_string(message).map_err(crate::Error::into_internal_error)
            }
            Self::Malformed { raw, .. } => Ok(raw.clone()),
        }
    }
}

/// Transport outgoing actor for line streams: serializes [`TransportFrame`] values and yields
/// lines.
///
/// This is a line-based variant of `transport_outgoing_actor` that works with a Sink<String>
/// instead of an AsyncWrite byte stream. This enables interception of lines before they are
/// written to the underlying transport.
///
/// This actor handles transport mechanics:
/// - Serializes single messages and malformed values to JSON strings
/// - Yields newline-terminated strings
/// - Handles serialization errors
///
/// This is the transport layer - it has no knowledge of protocol semantics (IDs, correlation, etc.).
async fn transport_outgoing_frames_actor(
    transport_rx: impl futures::Stream<Item = TransportFrame>,
    outgoing_lines: impl futures::Sink<String, Error = std::io::Error>,
) -> Result<(), crate::Error> {
    use futures::SinkExt;
    let mut transport_rx = pin!(transport_rx);
    let mut outgoing_lines = pin!(outgoing_lines);

    while let Some(frame) = transport_rx.next().await {
        let json_rpc_message = match frame {
            TransportFrame::Single(message) => message,
            TransportFrame::Malformed { raw, .. } => {
                let raw = malformed_line_value(raw)?;
                tracing::trace!(
                    message = %crate::util::redacted_json_line(&raw),
                    "Relaying invalid JSON-RPC value"
                );
                outgoing_lines
                    .send(raw)
                    .await
                    .map_err(crate::Error::into_internal_error)?;
                continue;
            }
        };
        match serde_json::to_string(&json_rpc_message) {
            Ok(line) => {
                tracing::trace!(message = %crate::util::redacted_json_line(&line), "Sending JSON-RPC message");
                outgoing_lines
                    .send(line)
                    .await
                    .map_err(crate::Error::into_internal_error)?;
            }

            Err(serialization_error) => {
                match json_rpc_message {
                    RawJsonRpcMessage::Request(_) | RawJsonRpcMessage::Notification(_) => {
                        // If we failed to serialize a request,
                        // just ignore it.
                        //
                        // Q: (Maybe it'd be nice to "reply" with an error?)
                        tracing::error!(
                            ?serialization_error,
                            "Failed to serialize request, ignoring"
                        );
                    }
                    RawJsonRpcMessage::Response(response) => {
                        // If we failed to serialize a *response*,
                        // send an error in response.
                        let id = match response {
                            Response::Result { id, .. } | Response::Error { id, .. } => id,
                        };
                        tracing::error!(
                            ?serialization_error,
                            ?id,
                            "Failed to serialize response, sending internal_error instead"
                        );
                        let error_line = serde_json::to_string(&RawJsonRpcMessage::response(
                            id,
                            Err(crate::Error::internal_error()),
                        ))
                        .unwrap();
                        outgoing_lines
                            .send(error_line)
                            .await
                            .map_err(crate::Error::into_internal_error)?;
                    }
                }
            }
        }
    }
    outgoing_lines
        .close()
        .await
        .map_err(crate::Error::into_internal_error)
}

/// A newline writer whose sink close reaches the actual write half. An unfold
/// sink can flush each line, but has no way to forward `AsyncWrite::close`.
pub(super) struct LineWriter<W> {
    writer: std::pin::Pin<Box<W>>,
    bytes: Vec<u8>,
    written: usize,
    closing: bool,
}

impl<W> LineWriter<W> {
    pub(super) fn new(writer: W) -> Self {
        Self {
            writer: Box::pin(writer),
            bytes: Vec::new(),
            written: 0,
            closing: false,
        }
    }
}

impl<W: futures::AsyncWrite> futures::Sink<String> for LineWriter<W> {
    type Error = std::io::Error;

    fn poll_ready(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        self.poll_flush(cx)
    }

    fn start_send(self: std::pin::Pin<&mut Self>, line: String) -> Result<(), Self::Error> {
        let this = self.get_mut();
        this.bytes = line.into_bytes();
        this.bytes.push(b'\n');
        this.written = 0;
        Ok(())
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        let this = self.get_mut();
        while this.written < this.bytes.len() {
            let count = futures::ready!(
                this.writer
                    .as_mut()
                    .poll_write(cx, &this.bytes[this.written..])
            )?;
            if count == 0 {
                return std::task::Poll::Ready(Err(std::io::ErrorKind::WriteZero.into()));
            }
            this.written += count;
        }
        this.bytes.clear();
        this.written = 0;
        this.writer.as_mut().poll_flush(cx)
    }

    fn poll_close(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), Self::Error>> {
        if !self.closing {
            futures::ready!(self.as_mut().poll_flush(cx))?;
            self.closing = true;
        }
        self.get_mut().writer.as_mut().poll_close(cx)
    }
}

fn malformed_line_value(raw: String) -> Result<String, crate::Error> {
    if !raw.contains('\r') && !raw.contains('\n') {
        return Ok(raw);
    }

    match serde_json::from_str::<serde_json::Value>(&raw) {
        Ok(value) => serde_json::to_string(&value),
        Err(_) => serde_json::to_string(&raw),
    }
    .map_err(crate::Error::into_internal_error)
}

pub(super) async fn transport_outgoing_lines_actor(
    transport_rx: impl futures::Stream<Item = TransportFrame>,
    outgoing_lines: impl futures::Sink<String, Error = std::io::Error>,
) -> Result<(), crate::Error> {
    transport_outgoing_frames_actor(transport_rx, outgoing_lines).await
}

/// Transport incoming actor for line streams: parses lines into [`TransportFrame`] values.
///
/// This is a line-based variant of `transport_incoming_actor` that works with a
/// Stream<Item = io::Result<String>> instead of an AsyncRead byte stream. This enables
/// interception of lines before they are parsed.
///
/// This actor handles transport mechanics:
/// - Reads lines from the stream
/// - Parses individual messages
/// - Retains malformed JSON, invalid messages, and unsupported batch arrays as
///   malformed frames
///
/// This is the transport layer - it has no knowledge of protocol semantics.
pub(super) async fn transport_incoming_lines_actor(
    incoming_lines: impl futures::Stream<Item = std::io::Result<String>>,
    transport_tx: mpsc::UnboundedSender<TransportFrame>,
) -> Result<(), crate::Error> {
    let mut incoming_lines = pin!(incoming_lines);
    while let Some(line_result) = incoming_lines.next().await {
        let line = line_result.map_err(crate::Error::into_internal_error)?;
        tracing::trace!(message = %crate::util::redacted_json_line(&line), "Received JSON-RPC message");

        transport_tx
            .unbounded_send(parse_incoming_line(&line))
            .map_err(crate::Error::into_internal_error)?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use super::*;
    use crate::ErrorCode;

    #[derive(Default)]
    struct PendingCloseWriter {
        bytes: Vec<u8>,
        close_polls: usize,
    }

    impl futures::AsyncWrite for PendingCloseWriter {
        fn poll_write(
            mut self: std::pin::Pin<&mut Self>,
            _: &mut std::task::Context<'_>,
            bytes: &[u8],
        ) -> std::task::Poll<std::io::Result<usize>> {
            self.bytes.extend_from_slice(bytes);
            std::task::Poll::Ready(Ok(bytes.len()))
        }

        fn poll_flush(
            self: std::pin::Pin<&mut Self>,
            _: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            assert_eq!(self.close_polls, 0, "do not flush again during shutdown");
            std::task::Poll::Ready(Ok(()))
        }

        fn poll_close(
            mut self: std::pin::Pin<&mut Self>,
            cx: &mut std::task::Context<'_>,
        ) -> std::task::Poll<std::io::Result<()>> {
            self.close_polls += 1;
            if self.close_polls == 1 {
                cx.waker().wake_by_ref();
                std::task::Poll::Pending
            } else {
                std::task::Poll::Ready(Ok(()))
            }
        }
    }

    #[test]
    fn byte_writer_flushes_then_continues_pending_shutdown_without_reflushing() {
        use futures::SinkExt as _;
        let mut sink = LineWriter::new(PendingCloseWriter::default());
        futures::executor::block_on(async {
            sink.send("line".to_string()).await.unwrap();
            sink.close().await.unwrap();
        });
        assert_eq!(sink.writer.bytes, b"line\n");
        assert_eq!(sink.writer.close_polls, 2);
    }

    #[test]
    fn preserves_malformed_standalone_response() {
        let TransportFrame::Malformed { error, .. } = parse_incoming_line(
            r#"{"jsonrpc":"2.0","id":1,"result":null,"error":{"code":-32603,"message":"Internal error"}}"#,
        ) else {
            panic!("expected one retained invalid response");
        };

        assert_eq!(error.code, ErrorCode::INVALID_REQUEST);
    }

    #[test]
    fn preserves_malformed_call_shaped_standalone_message() {
        let TransportFrame::Malformed { error, .. } =
            parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"method":"one","result":null}"#)
        else {
            panic!("expected one invalid-request error");
        };

        assert_eq!(error.code, ErrorCode::INVALID_REQUEST);
    }

    #[test]
    fn parses_valid_standalone_response() {
        assert!(matches!(
            parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"result":{"ok":true}}"#),
            TransportFrame::Single(RawJsonRpcMessage::Response(_))
        ));
    }

    #[test]
    fn batches_are_one_invalid_request() {
        for line in [
            "[]",
            "[1, 2, 3]",
            r#"[{"jsonrpc":"2.0","id":1,"method":"one","params":{}}]"#,
            r#"[{"jsonrpc":"2.0","id":1,"result":null}]"#,
        ] {
            let TransportFrame::Malformed { raw, error } = parse_incoming_line(line) else {
                panic!("expected one invalid-request error for {line}");
            };
            assert_eq!(raw, line);
            assert_eq!(error.code, ErrorCode::INVALID_REQUEST);
        }
    }

    #[test]
    fn malformed_json_is_a_parse_error() {
        let TransportFrame::Malformed { raw, error } = parse_incoming_line("[") else {
            panic!("expected one parse error");
        };

        assert_eq!(raw, "[");
        assert_eq!(error.code, ErrorCode::PARSE_ERROR);
    }

    #[test]
    fn valid_json_with_an_invalid_envelope_is_an_invalid_request() {
        let TransportFrame::Malformed { raw, error } = parse_incoming_line("17") else {
            panic!("expected one invalid-request error");
        };

        assert_eq!(raw, "17");
        assert_eq!(error.code, ErrorCode::INVALID_REQUEST);
    }

    #[tokio::test]
    async fn multiline_malformed_frame_is_written_as_one_line_value() {
        let raw = "not json\r\n{\"jsonrpc\":\"2.0\",\"method\":\"injected\"}".to_string();
        let captured = Arc::new(Mutex::new(Vec::new()));
        let outgoing = futures::sink::unfold(captured.clone(), |captured, line| async move {
            captured.lock().unwrap().push(line);
            Ok::<_, std::io::Error>(captured)
        });

        transport_outgoing_frames_actor(
            futures::stream::iter([TransportFrame::Malformed {
                raw: raw.clone(),
                error: crate::Error::parse_error(),
            }]),
            outgoing,
        )
        .await
        .unwrap();

        let lines = captured.lock().unwrap();
        assert_eq!(lines.len(), 1);
        assert!(!lines[0].contains('\r') && !lines[0].contains('\n'));
        assert_eq!(serde_json::from_str::<String>(&lines[0]).unwrap(), raw);
    }

    #[test]
    fn multiline_invalid_json_rpc_value_is_compacted_without_changing_value() {
        let raw = "{\n  \"jsonrpc\": \"2.0\",\n  \"method\": 1\n}".to_string();
        let expected = serde_json::from_str::<serde_json::Value>(&raw).unwrap();
        let line = malformed_line_value(raw).unwrap();

        assert!(!line.contains('\r') && !line.contains('\n'));
        assert_eq!(
            serde_json::from_str::<serde_json::Value>(&line).unwrap(),
            expected
        );
    }
}