deepstrike_sdk/runtime/
process_sandbox_plane.rs1use 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 pub sandbox_dir: PathBuf,
19 pub allowed_env_keys: Vec<String>,
21 pub timeout_ms: u64,
23 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
38pub 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
83async 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 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 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
186fn 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}