orra 0.0.2

Context-aware agent session management for any application
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

use async_trait::async_trait;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};

use super::types::{JsonRpcRequest, JsonRpcResponse};

#[derive(Debug, thiserror::Error)]
pub enum TransportError {
    #[error("io error: {0}")]
    Io(String),

    #[error("serialization error: {0}")]
    Serialization(String),

    #[error("transport closed")]
    Closed,
}

/// Transport layer for communicating with an MCP server.
#[async_trait]
pub trait McpTransport: Send + Sync {
    async fn send(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse, TransportError>;
    async fn close(&self) -> Result<(), TransportError>;
}

/// Stdio transport: communicates with a child process via stdin/stdout JSON lines.
pub struct StdioTransport {
    stdin: tokio::sync::Mutex<tokio::process::ChildStdin>,
    stdout: tokio::sync::Mutex<BufReader<tokio::process::ChildStdout>>,
    child: tokio::sync::Mutex<Child>,
}

impl StdioTransport {
    /// Spawn a child process and create a transport for communicating with it.
    pub async fn spawn(program: &str, args: &[&str]) -> Result<Self, TransportError> {
        Self::spawn_with_env(program, args, &[]).await
    }

    /// Spawn a child process with extra environment variables.
    pub async fn spawn_with_env(
        program: &str,
        args: &[&str],
        env: &[(&str, &str)],
    ) -> Result<Self, TransportError> {
        let mut cmd = Command::new(program);
        cmd.args(args)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::null());

        for (key, value) in env {
            cmd.env(key, value);
        }

        let mut child = cmd
            .spawn()
            .map_err(|e| TransportError::Io(format!("failed to spawn {}: {}", program, e)))?;

        let stdin = child.stdin.take()
            .ok_or_else(|| TransportError::Io("no stdin".into()))?;
        let stdout = child.stdout.take()
            .ok_or_else(|| TransportError::Io("no stdout".into()))?;

        Ok(Self {
            stdin: tokio::sync::Mutex::new(stdin),
            stdout: tokio::sync::Mutex::new(BufReader::new(stdout)),
            child: tokio::sync::Mutex::new(child),
        })
    }
}

#[async_trait]
impl McpTransport for StdioTransport {
    async fn send(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse, TransportError> {
        let mut json = serde_json::to_string(&request)
            .map_err(|e| TransportError::Serialization(e.to_string()))?;
        json.push('\n');

        let mut stdin = self.stdin.lock().await;
        stdin.write_all(json.as_bytes()).await
            .map_err(|e| TransportError::Io(e.to_string()))?;
        stdin.flush().await
            .map_err(|e| TransportError::Io(e.to_string()))?;

        let mut stdout = self.stdout.lock().await;
        let mut line = String::new();
        let bytes = stdout.read_line(&mut line).await
            .map_err(|e| TransportError::Io(e.to_string()))?;

        if bytes == 0 {
            return Err(TransportError::Closed);
        }

        let response: JsonRpcResponse = serde_json::from_str(line.trim())
            .map_err(|e| TransportError::Serialization(format!("parse response: {}", e)))?;

        Ok(response)
    }

    async fn close(&self) -> Result<(), TransportError> {
        let mut child = self.child.lock().await;
        let _ = child.kill().await;
        Ok(())
    }
}

/// In-memory transport for testing. Uses mpsc channels to simulate a server.
#[cfg(test)]
pub(crate) struct MockTransport {
    tx: tokio::sync::mpsc::Sender<JsonRpcRequest>,
    rx: tokio::sync::Mutex<tokio::sync::mpsc::Receiver<JsonRpcResponse>>,
}

#[cfg(test)]
impl MockTransport {
    pub fn new() -> (
        Self,
        tokio::sync::mpsc::Receiver<JsonRpcRequest>,
        tokio::sync::mpsc::Sender<JsonRpcResponse>,
    ) {
        let (req_tx, req_rx) = tokio::sync::mpsc::channel(16);
        let (resp_tx, resp_rx) = tokio::sync::mpsc::channel(16);
        let transport = Self {
            tx: req_tx,
            rx: tokio::sync::Mutex::new(resp_rx),
        };
        (transport, req_rx, resp_tx)
    }
}

#[cfg(test)]
#[async_trait]
impl McpTransport for MockTransport {
    async fn send(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse, TransportError> {
        self.tx.send(request).await
            .map_err(|_| TransportError::Closed)?;
        let mut rx = self.rx.lock().await;
        rx.recv().await.ok_or(TransportError::Closed)
    }

    async fn close(&self) -> Result<(), TransportError> {
        Ok(())
    }
}