Skip to main content

deepstrike_sdk/runtime/
process_sandbox_plane.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use deepstrike_core::types::message::{ToolCall, ToolSchema};
7use futures::stream::Stream;
8use serde_json::json;
9use tokio::io::AsyncReadExt;
10
11use crate::Result;
12use crate::run_event::RunEvent;
13use crate::runtime::execution_plane::{ExecutionPlane, LocalExecutionPlane, RunContext};
14use crate::tools::RegisteredTool;
15
16pub struct SandboxOptions {
17    /// Working directory for all subprocesses; isolated from the host file system by convention.
18    pub sandbox_dir: PathBuf,
19    /// Host env-var names forwarded into subprocesses. Default: none.
20    pub allowed_env_keys: Vec<String>,
21    /// Per-call hard timeout in ms. Default: 30 000.
22    pub timeout_ms: u64,
23    /// Truncate stdout+stderr after this many bytes. Default: 1 MiB.
24    pub max_output_bytes: usize,
25}
26
27impl Default for SandboxOptions {
28    fn default() -> Self {
29        Self {
30            sandbox_dir: std::env::temp_dir().join("deepstrike-sandbox"),
31            allowed_env_keys: vec![],
32            timeout_ms: 30_000,
33            max_output_bytes: 1_048_576,
34        }
35    }
36}
37
38/// `LocalExecutionPlane` extended with two subprocess built-in tools:
39///   - `run_bash`   — executes a bash command inside `sandbox_dir`.
40///   - `run_python` — evaluates a Python script inside `sandbox_dir`.
41///
42/// All registered Rust tools still run in-process (identical to `LocalExecutionPlane`).
43/// Subprocesses are launched with `sandbox_dir` as cwd and a stripped environment.
44/// This is an execution hygiene boundary, not an OS-enforced filesystem sandbox.
45pub struct ProcessSandboxPlane {
46    inner: LocalExecutionPlane,
47}
48
49impl ProcessSandboxPlane {
50    pub fn new(opts: SandboxOptions) -> Self {
51        let opts = Arc::new(opts);
52        let mut plane = LocalExecutionPlane::new();
53        plane.register(make_bash_tool(Arc::clone(&opts)));
54        plane.register(make_python_tool(opts));
55        Self { inner: plane }
56    }
57
58    pub fn register(&mut self, tool: RegisteredTool) -> &mut Self {
59        self.inner.register(tool);
60        self
61    }
62
63    pub fn unregister(&mut self, name: &str) -> &mut Self {
64        self.inner.unregister(name);
65        self
66    }
67}
68
69impl ExecutionPlane for ProcessSandboxPlane {
70    fn schemas(&self) -> Vec<ToolSchema> {
71        self.inner.schemas()
72    }
73
74    fn execute_all<'a>(
75        &'a self,
76        calls: &'a [ToolCall],
77        ctx: RunContext<'a>,
78    ) -> Pin<Box<dyn Stream<Item = Result<RunEvent>> + Send + 'a>> {
79        self.inner.execute_all(calls, ctx)
80    }
81}
82
83// ── Subprocess runner ─────────────────────────────────────────────────────────
84
85async fn run_subprocess(
86    cmd: &'static str,
87    script: String,
88    sandbox_dir: PathBuf,
89    allowed_env_keys: Vec<String>,
90    timeout_ms: u64,
91    max_output_bytes: usize,
92) -> (String, bool) {
93    if let Err(e) = tokio::fs::create_dir_all(&sandbox_dir).await {
94        return (e.to_string(), true);
95    }
96
97    let mut env: HashMap<String, String> = HashMap::new();
98    let dir = sandbox_dir.to_string_lossy().to_string();
99    env.insert("HOME".into(), dir.clone());
100    env.insert("TMPDIR".into(), dir);
101    env.insert("PATH".into(), "/usr/local/bin:/usr/bin:/bin".into());
102    for key in &allowed_env_keys {
103        if let Ok(val) = std::env::var(key) {
104            env.insert(key.clone(), val);
105        }
106    }
107
108    let mut child = match tokio::process::Command::new(cmd)
109        .arg("-c")
110        .arg(&script)
111        .current_dir(&sandbox_dir)
112        .envs(&env)
113        .stdin(std::process::Stdio::null())
114        .stdout(std::process::Stdio::piped())
115        .stderr(std::process::Stdio::piped())
116        .spawn()
117    {
118        Ok(c) => c,
119        Err(e) => return (e.to_string(), true),
120    };
121
122    // Drain stdout/stderr concurrently so the child never blocks on a full pipe buffer.
123    let stdout_pipe = child.stdout.take().expect("stdout was piped");
124    let stderr_pipe = child.stderr.take().expect("stderr was piped");
125    let read_out = tokio::spawn(async move {
126        let mut buf = Vec::new();
127        tokio::io::BufReader::new(stdout_pipe)
128            .read_to_end(&mut buf)
129            .await
130            .ok();
131        buf
132    });
133    let read_err = tokio::spawn(async move {
134        let mut buf = Vec::new();
135        tokio::io::BufReader::new(stderr_pipe)
136            .read_to_end(&mut buf)
137            .await
138            .ok();
139        buf
140    });
141
142    // wait() takes &mut self, so child remains usable if the timeout fires.
143    let timed_out =
144        tokio::time::timeout(tokio::time::Duration::from_millis(timeout_ms), child.wait()).await;
145
146    let is_error = match timed_out {
147        Ok(Ok(status)) => !status.success(),
148        Ok(Err(e)) => {
149            read_out.abort();
150            read_err.abort();
151            return (e.to_string(), true);
152        }
153        Err(_) => {
154            child.kill().await.ok();
155            child.wait().await.ok();
156            read_out.abort();
157            read_err.abort();
158            return (format!("timed out after {timeout_ms}ms"), true);
159        }
160    };
161
162    let out_bytes = read_out.await.unwrap_or_default();
163    let err_bytes = read_err.await.unwrap_or_default();
164    let mut combined = [out_bytes, err_bytes].concat();
165    if combined.len() > max_output_bytes {
166        combined.truncate(max_output_bytes);
167        combined.extend_from_slice(b"\n[output truncated]");
168    }
169    let text = String::from_utf8_lossy(&combined).into_owned();
170    if is_error && text.trim().is_empty() {
171        return (
172            "Process exited with non-zero status and produced no output.".into(),
173            true,
174        );
175    }
176    (
177        if text.is_empty() {
178            "(no output)".into()
179        } else {
180            text
181        },
182        is_error,
183    )
184}
185
186// ── Tool constructors ─────────────────────────────────────────────────────────
187
188fn make_bash_tool(opts: Arc<SandboxOptions>) -> RegisteredTool {
189    RegisteredTool::text(
190        "run_bash",
191        "Run a bash command with the sandbox directory as cwd and a stripped environment. \
192         This is not an OS-enforced filesystem sandbox.",
193        json!({
194            "type": "object",
195            "properties": {
196                "command": { "type": "string", "description": "The bash command to execute." }
197            },
198            "required": ["command"]
199        }),
200        move |args| {
201            let opts = Arc::clone(&opts);
202            Box::pin(async move {
203                let command = args["command"].as_str().unwrap_or("").to_string();
204                let (output, _) = run_subprocess(
205                    "bash",
206                    command,
207                    opts.sandbox_dir.clone(),
208                    opts.allowed_env_keys.clone(),
209                    opts.timeout_ms,
210                    opts.max_output_bytes,
211                )
212                .await;
213                Ok(output)
214            })
215        },
216    )
217}
218
219fn make_python_tool(opts: Arc<SandboxOptions>) -> RegisteredTool {
220    RegisteredTool::text(
221        "run_python",
222        "Evaluate a Python script with the sandbox directory as cwd and a stripped environment.",
223        json!({
224            "type": "object",
225            "properties": {
226                "code": { "type": "string", "description": "The Python code to evaluate." }
227            },
228            "required": ["code"]
229        }),
230        move |args| {
231            let opts = Arc::clone(&opts);
232            Box::pin(async move {
233                let code = args["code"].as_str().unwrap_or("").to_string();
234                let (output, _) = run_subprocess(
235                    "python3",
236                    code,
237                    opts.sandbox_dir.clone(),
238                    opts.allowed_env_keys.clone(),
239                    opts.timeout_ms,
240                    opts.max_output_bytes,
241                )
242                .await;
243                Ok(output)
244            })
245        },
246    )
247}