IdeA/crates/backend/src/mcp_bridge.rs

//! The `idea mcp-server` bridge shared by desktop and standalone web binaries.
//!
//! An MCP CLI spawns `<exe> mcp-server --endpoint <…> --project <…> --requester <…>`.
//! This process stays headless and relays JSON-RPC lines between stdio and the
//! per-project local loopback endpoint where the real MCP server listens.

use std::process::ExitCode;
use std::time::Duration;

use interprocess::local_socket::tokio::Stream as LocalSocketStream;
use interprocess::local_socket::traits::tokio::Stream as _;
use interprocess::local_socket::{GenericFilePath, ToFsName as _};
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};

const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
const DRAIN_GRACE: Duration = Duration::from_secs(1);

/// Parsed `mcp-server` invocation arguments.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BridgeArgs {
    /// The loopback address to connect to.
    pub endpoint: String,
    /// The project id forwarded in the handshake.
    pub project: String,
    /// The requester id forwarded in the handshake.
    pub requester: String,
}

impl BridgeArgs {
    /// Parses `mcp-server`'s own arguments.
    pub fn parse<I, S>(args: I) -> Result<Self, String>
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let mut endpoint: Option<String> = None;
        let mut project = String::new();
        let mut requester = String::new();

        let mut it = args.into_iter().map(Into::into);
        while let Some(flag) = it.next() {
            match flag.as_str() {
                "--endpoint" => {
                    endpoint = Some(
                        it.next()
                            .ok_or_else(|| "--endpoint requires a value".to_string())?,
                    );
                }
                "--project" => {
                    project = it
                        .next()
                        .ok_or_else(|| "--project requires a value".to_string())?;
                }
                "--requester" => {
                    requester = it
                        .next()
                        .ok_or_else(|| "--requester requires a value".to_string())?;
                }
                other => return Err(format!("unknown argument: {other}")),
            }
        }

        let endpoint = endpoint.ok_or_else(|| "--endpoint is required".to_string())?;
        Ok(Self {
            endpoint,
            project,
            requester,
        })
    }

    fn handshake_line(&self) -> Vec<u8> {
        let line = serde_json::json!({
            "project": self.project,
            "requester": self.requester,
        })
        .to_string();
        let mut bytes = line.into_bytes();
        bytes.push(b'\n');
        bytes
    }
}

/// Synchronous entry point for the headless stdio bridge.
#[must_use]
pub fn run_mcp_bridge<I, S>(args: I) -> ExitCode
where
    I: IntoIterator<Item = S>,
    S: Into<String>,
{
    let args = match BridgeArgs::parse(args) {
        Ok(a) => a,
        Err(e) => {
            eprintln!("idea mcp-server: {e}");
            return ExitCode::FAILURE;
        }
    };

    let runtime = match tokio::runtime::Builder::new_current_thread()
        .enable_io()
        .enable_time()
        .build()
    {
        Ok(rt) => rt,
        Err(e) => {
            eprintln!("idea mcp-server: failed to start runtime: {e}");
            return ExitCode::FAILURE;
        }
    };

    runtime.block_on(async move {
        match bridge_over_loopback(&args).await {
            Ok(()) => ExitCode::SUCCESS,
            Err(e) => {
                eprintln!("idea mcp-server: {e}");
                ExitCode::FAILURE
            }
        }
    })
}

async fn bridge_over_loopback(args: &BridgeArgs) -> Result<(), String> {
    let conn = connect_loopback(&args.endpoint).await?;
    let (lp_read, lp_write) = tokio::io::split(conn);
    relay(
        args,
        tokio::io::stdin(),
        tokio::io::stdout(),
        lp_read,
        lp_write,
    )
    .await
}

async fn connect_loopback(endpoint: &str) -> Result<LocalSocketStream, String> {
    let name = endpoint
        .to_fs_name::<GenericFilePath>()
        .map_err(|e| format!("invalid endpoint {endpoint:?}: {e}"))?;

    match tokio::time::timeout(CONNECT_TIMEOUT, LocalSocketStream::connect(name)).await {
        Ok(Ok(stream)) => Ok(stream),
        Ok(Err(e)) => Err(format!("cannot connect to endpoint {endpoint:?}: {e}")),
        Err(_) => Err(format!(
            "timed out connecting to endpoint {endpoint:?} after {}s",
            CONNECT_TIMEOUT.as_secs()
        )),
    }
}

async fn relay<CIn, COut, LIn, LOut>(
    args: &BridgeArgs,
    cli_in: CIn,
    cli_out: COut,
    lp_read: LIn,
    mut lp_write: LOut,
) -> Result<(), String>
where
    CIn: AsyncRead + Unpin,
    COut: AsyncWrite + Unpin,
    LIn: AsyncRead + Unpin,
    LOut: AsyncWrite + Unpin,
{
    lp_write
        .write_all(&args.handshake_line())
        .await
        .map_err(|e| format!("handshake write failed: {e}"))?;
    lp_write
        .flush()
        .await
        .map_err(|e| format!("handshake flush failed: {e}"))?;

    let cli_to_lp = pump_lines(BufReader::new(cli_in), lp_write, "stdin", "loopback");
    let lp_to_cli = pump_lines(BufReader::new(lp_read), cli_out, "loopback", "stdout");
    tokio::pin!(cli_to_lp);
    tokio::pin!(lp_to_cli);

    tokio::select! {
        r = &mut cli_to_lp => {
            let mut lp_write = r?;
            let _ = lp_write.shutdown().await;
            let _ = tokio::time::timeout(DRAIN_GRACE, &mut lp_to_cli).await;
            Ok(())
        }
        r = &mut lp_to_cli => r.map(|_| ()),
    }
}

async fn pump_lines<R, W>(
    mut reader: BufReader<R>,
    mut writer: W,
    src: &str,
    dst: &str,
) -> Result<W, String>
where
    R: AsyncRead + Unpin,
    W: AsyncWrite + Unpin,
{
    let mut line = String::new();
    loop {
        line.clear();
        let n = reader
            .read_line(&mut line)
            .await
            .map_err(|e| format!("{src} read failed: {e}"))?;
        if n == 0 {
            writer
                .flush()
                .await
                .map_err(|e| format!("{dst} flush failed: {e}"))?;
            return Ok(writer);
        }
        writer
            .write_all(line.as_bytes())
            .await
            .map_err(|e| format!("{dst} write failed: {e}"))?;
        writer
            .flush()
            .await
            .map_err(|e| format!("{dst} flush failed: {e}"))?;
    }
}

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

    #[test]
    fn parses_all_three_flags() {
        let args = BridgeArgs::parse([
            "--endpoint",
            "/tmp/x.sock",
            "--project",
            "p1",
            "--requester",
            "agent-7",
        ])
        .unwrap();
        assert_eq!(args.endpoint, "/tmp/x.sock");
        assert_eq!(args.project, "p1");
        assert_eq!(args.requester, "agent-7");
    }

    #[test]
    fn handshake_line_carries_identity_as_json() {
        let args = BridgeArgs {
            endpoint: "/tmp/x.sock".into(),
            project: "proj".into(),
            requester: "rq".into(),
        };
        let line = args.handshake_line();
        assert_eq!(*line.last().unwrap(), b'\n');
        let value: serde_json::Value = serde_json::from_slice(&line[..line.len() - 1]).unwrap();
        assert_eq!(value["project"], "proj");
        assert_eq!(value["requester"], "rq");
    }

    #[test]
    fn missing_endpoint_returns_failure() {
        assert_eq!(run_mcp_bridge(["--project", "p"]), ExitCode::FAILURE);
    }
}