chimera_core/transport/
subprocess.rs

1use std::collections::{HashMap, HashSet};
2use std::path::PathBuf;
3use std::process::ExitStatus;
4
5use serde::Serialize;
6use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
7use tokio::process::{Child, ChildStdin, ChildStdout, Command};
8use tokio::task::JoinHandle;
9
10use crate::AgentError;
11
12const MAX_STDERR_SIZE: usize = 64 * 1024; // 64 KB
13
14#[derive(Debug, Clone, bon::Builder)]
15pub struct SubprocessCommand {
16    pub binary: PathBuf,
17    #[builder(default)]
18    pub args: Vec<String>,
19    #[builder(default)]
20    pub env: HashMap<String, String>,
21    /// Inherited environment variables to remove from the subprocess.
22    /// Applied after `env` overlay; only meaningful when `inherit_env` is true.
23    #[builder(default)]
24    pub env_remove: HashSet<String>,
25    /// If false, the child gets ONLY `env`. If true (default), inherits
26    /// the current process env and overlays `env`.
27    #[builder(default = true)]
28    pub inherit_env: bool,
29    pub cwd: Option<PathBuf>,
30}
31
32pub struct Subprocess {
33    child: Child,
34    stdin: Option<BufWriter<ChildStdin>>,
35    stdout: Option<ChildStdout>,
36    stderr_task: JoinHandle<String>,
37}
38
39impl Subprocess {
40    pub fn spawn(cmd: SubprocessCommand) -> Result<Self, AgentError> {
41        let mut command = Command::new(&cmd.binary);
42        command.args(&cmd.args);
43
44        if cmd.inherit_env {
45            command.envs(&cmd.env);
46            for key in &cmd.env_remove {
47                command.env_remove(key);
48            }
49        } else {
50            command.env_clear().envs(&cmd.env);
51        }
52
53        if let Some(cwd) = &cmd.cwd {
54            command.current_dir(cwd);
55        }
56
57        command
58            .stdin(std::process::Stdio::piped())
59            .stdout(std::process::Stdio::piped())
60            .stderr(std::process::Stdio::piped());
61        command.kill_on_drop(true);
62
63        let mut child = command
64            .spawn()
65            .map_err(|source| AgentError::SpawnFailed { source })?;
66
67        let stdin = child.stdin.take().map(BufWriter::new);
68        let stdout = child.stdout.take();
69
70        let stderr = child.stderr.take().expect("stderr was piped");
71        let stderr_task = tokio::spawn(async move {
72            let mut buf = vec![0u8; MAX_STDERR_SIZE];
73            let mut reader = BufReader::new(stderr);
74            let mut total = 0;
75
76            loop {
77                match reader.read(&mut buf[total..]).await {
78                    Ok(0) => break,
79                    Ok(n) => {
80                        total += n;
81                        if total >= MAX_STDERR_SIZE {
82                            break;
83                        }
84                    }
85                    Err(_) => break,
86                }
87            }
88
89            String::from_utf8_lossy(&buf[..total]).into_owned()
90        });
91
92        Ok(Self {
93            child,
94            stdin,
95            stdout,
96            stderr_task,
97        })
98    }
99
100    pub fn take_stdin(&mut self) -> Option<BufWriter<ChildStdin>> {
101        self.stdin.take()
102    }
103
104    pub fn take_stdout(&mut self) -> Option<ChildStdout> {
105        self.stdout.take()
106    }
107
108    pub async fn write(&mut self, data: &str) -> Result<(), AgentError> {
109        let stdin = self.stdin.as_mut().ok_or_else(|| AgentError::Other {
110            message: "stdin already closed".into(),
111            source: None,
112        })?;
113        stdin
114            .write_all(data.as_bytes())
115            .await
116            .map_err(|e| AgentError::StdinWrite { source: e })?;
117        stdin
118            .flush()
119            .await
120            .map_err(|e| AgentError::StdinWrite { source: e })
121    }
122
123    pub async fn write_json(&mut self, value: &impl Serialize) -> Result<(), AgentError> {
124        let json = serde_json::to_string(value).map_err(|e| AgentError::Other {
125            message: "failed to serialize JSON".into(),
126            source: Some(Box::new(e)),
127        })?;
128        self.write(&json).await?;
129        self.write("\n").await
130    }
131
132    pub async fn close_stdin(&mut self) {
133        self.stdin.take();
134    }
135
136    pub async fn kill(&mut self) -> Result<(), AgentError> {
137        self.child.kill().await.map_err(|e| AgentError::Other {
138            message: "failed to kill process".into(),
139            source: Some(Box::new(e)),
140        })
141    }
142
143    pub async fn wait(mut self) -> Result<SubprocessOutput, AgentError> {
144        let status = self.child.wait().await.map_err(|e| AgentError::Other {
145            message: "failed to wait for process".into(),
146            source: Some(Box::new(e)),
147        })?;
148        let stderr = self.stderr_task.await.unwrap_or_default();
149        Ok(SubprocessOutput { status, stderr })
150    }
151}
152
153#[derive(Debug)]
154pub struct SubprocessOutput {
155    pub status: ExitStatus,
156    pub stderr: String,
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn stderr_cap_threshold() {
165        assert_eq!(MAX_STDERR_SIZE, 64 * 1024);
166    }
167}
chimera_core/transport/subprocess.rs

chimera_core/transport/
subprocess.rs
