IdeaSDK/crates/infrastructure/src/orchestrator/mcp/transport.rs

//! Concrete transports for the MCP adapter.
//!
//! - [`StdioTransport`] — the **default** transport (cadrage S-MCP): newline-
//!   delimited JSON ("JSON Lines") over a pair of async byte streams. A CLI that
//!   IdeA spawns with the injected MCP config speaks to the server over its
//!   stdin/stdout; IdeA bridges those streams here.
//! - [`MemoryTransport`] — an in-memory, **scriptable** transport used by tests to
//!   drive the server with **no socket and no child process** (S-MCP testability
//!   requirement). Lives in production code (not `#[cfg(test)]`) so the test crate
//!   and downstream adapters can construct it.
//!
//! A `socket` transport is a documented **TODO** (cadrage S-MCP, point ouvert):
//! the [`Transport`] seam means it can be added without touching the server.

use std::collections::VecDeque;

use tokio::io::{AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::sync::mpsc;

use super::jsonrpc::{Transport, TransportError};

/// Newline-delimited JSON transport over a generic async reader/writer.
///
/// Generic over the streams so the default real wiring uses
/// `tokio::io::Stdin`/`Stdout`, while tests (or a future socket) can plug any
/// `AsyncRead`/`AsyncWrite`. Each `recv` reads one line; each `send` writes one
/// JSON value followed by `\n` and flushes.
pub struct StdioTransport<R, W> {
    reader: BufReader<R>,
    writer: W,
}

impl<R, W> StdioTransport<R, W>
where
    R: tokio::io::AsyncRead + Unpin + Send,
    W: AsyncWrite + Unpin + Send,
{
    /// Wraps a reader/writer pair as a line-delimited JSON transport.
    pub fn new(reader: R, writer: W) -> Self {
        Self {
            reader: BufReader::new(reader),
            writer,
        }
    }

    /// Wraps an **already-buffered** reader plus a writer as a line-delimited JSON
    /// transport.
    ///
    /// Used by the per-peer accept loop (M5c): the handshake line (cadrage v5 §1.4)
    /// is read off a `BufReader` over the loopback's read half *before* serving, and
    /// that same `BufReader` — carrying any bytes already buffered past the handshake
    /// — is handed here so no JSON-RPC byte is lost between the handshake and the
    /// serve loop.
    pub fn from_buffered(reader: BufReader<R>, writer: W) -> Self {
        Self { reader, writer }
    }
}

#[async_trait::async_trait]
impl<R, W> Transport for StdioTransport<R, W>
where
    R: tokio::io::AsyncRead + Unpin + Send,
    W: AsyncWrite + Unpin + Send,
{
    async fn recv(&mut self) -> Result<Vec<u8>, TransportError> {
        let mut line = String::new();
        let n = self
            .reader
            .read_line(&mut line)
            .await
            .map_err(|e| TransportError::Io(e.to_string()))?;
        if n == 0 {
            return Err(TransportError::Closed);
        }
        Ok(line.into_bytes())
    }

    async fn send(&mut self, message: &[u8]) -> Result<(), TransportError> {
        self.writer
            .write_all(message)
            .await
            .map_err(|e| TransportError::Io(e.to_string()))?;
        self.writer
            .write_all(b"\n")
            .await
            .map_err(|e| TransportError::Io(e.to_string()))?;
        self.writer
            .flush()
            .await
            .map_err(|e| TransportError::Io(e.to_string()))?;
        Ok(())
    }
}

/// An in-memory, scriptable transport for tests (zero I/O, zero network).
///
/// `recv` drains a pre-loaded queue of inbound messages and then reports
/// [`TransportError::Closed`] (so the serve loop terminates deterministically);
/// every `send` is captured on an `mpsc` channel the test can drain to assert the
/// exact responses the server produced.
pub struct MemoryTransport {
    inbound: VecDeque<Vec<u8>>,
    outbound: mpsc::UnboundedSender<Vec<u8>>,
}

impl MemoryTransport {
    /// Builds a transport that will hand `messages` to the server in order, then
    /// signal EOF. Returns the transport and a receiver of everything the server
    /// `send`s back.
    #[must_use]
    pub fn scripted(messages: Vec<Vec<u8>>) -> (Self, mpsc::UnboundedReceiver<Vec<u8>>) {
        let (tx, rx) = mpsc::unbounded_channel();
        (
            Self {
                inbound: messages.into(),
                outbound: tx,
            },
            rx,
        )
    }

    /// Convenience: script from JSON strings.
    #[must_use]
    pub fn scripted_str(messages: &[&str]) -> (Self, mpsc::UnboundedReceiver<Vec<u8>>) {
        Self::scripted(messages.iter().map(|m| m.as_bytes().to_vec()).collect())
    }
}

#[async_trait::async_trait]
impl Transport for MemoryTransport {
    async fn recv(&mut self) -> Result<Vec<u8>, TransportError> {
        self.inbound.pop_front().ok_or(TransportError::Closed)
    }

    async fn send(&mut self, message: &[u8]) -> Result<(), TransportError> {
        self.outbound
            .send(message.to_vec())
            .map_err(|_| TransportError::Closed)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn memory_transport_replays_then_closes() {
        let (mut t, _rx) = MemoryTransport::scripted_str(&["a", "b"]);
        assert_eq!(t.recv().await.unwrap(), b"a");
        assert_eq!(t.recv().await.unwrap(), b"b");
        assert!(matches!(t.recv().await, Err(TransportError::Closed)));
    }

    #[tokio::test]
    async fn memory_transport_captures_sends() {
        let (mut t, mut rx) = MemoryTransport::scripted(vec![]);
        t.send(b"hello").await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), b"hello");
    }

    #[tokio::test]
    async fn stdio_transport_round_trips_lines() {
        let input = b"{\"x\":1}\n{\"y\":2}\n".to_vec();
        let mut out: Vec<u8> = Vec::new();
        {
            let mut t = StdioTransport::new(&input[..], &mut out);
            assert_eq!(t.recv().await.unwrap(), b"{\"x\":1}\n");
            t.send(b"{\"ok\":true}").await.unwrap();
        }
        assert_eq!(out, b"{\"ok\":true}\n");
    }
}