1use std::sync::Arc;
6
7use base64::Engine;
8use serde_json::{Map, Value};
9use tokio::sync::mpsc;
10
11use crate::channel::{AsyncFrame, ControlChannel};
12use crate::client::Sandbox;
13use crate::error::{is_route_unavailable, SolariError};
14use crate::http::encode_uri_component;
15use crate::types::CommandResult;
16
17pub type OutputCallback = Arc<dyn Fn(&str) + Send + Sync>;
19
20#[derive(Default, Clone)]
22pub struct RunOptions {
23 pub args: Option<Vec<String>>,
25 pub cwd: Option<String>,
26 pub env: Option<std::collections::HashMap<String, String>>,
27 pub user: Option<String>,
28 pub timeout_ms: Option<u64>,
29 pub background: bool,
31 pub on_stdout: Option<OutputCallback>,
32 pub on_stderr: Option<OutputCallback>,
33}
34
35impl RunOptions {
36 pub fn new() -> Self {
37 Self::default()
38 }
39 pub fn args<I, S>(mut self, args: I) -> Self
40 where
41 I: IntoIterator<Item = S>,
42 S: Into<String>,
43 {
44 self.args = Some(args.into_iter().map(Into::into).collect());
45 self
46 }
47 pub fn cwd(mut self, cwd: impl Into<String>) -> Self {
48 self.cwd = Some(cwd.into());
49 self
50 }
51}
52
53pub struct CommandHandle {
55 pub cmd_id: String,
56 channel: Arc<ControlChannel>,
57 rx: mpsc::UnboundedReceiver<AsyncFrame>,
58 on_stdout: Option<OutputCallback>,
59 on_stderr: Option<OutputCallback>,
60}
61
62impl CommandHandle {
63 pub async fn stdin(&self, data: impl AsRef<[u8]>) -> Result<(), SolariError> {
65 let b64 = base64::engine::general_purpose::STANDARD.encode(data.as_ref());
66 let mut p = Map::new();
67 p.insert("cmdId".into(), Value::String(self.cmd_id.clone()));
68 p.insert("base64".into(), Value::String(b64));
69 self.channel.call("cmd.stdin", Value::Object(p)).await.map(|_| ())
70 }
71
72 pub async fn kill(&self, signal: Option<i64>) -> Result<(), SolariError> {
74 let mut p = Map::new();
75 p.insert("cmdId".into(), Value::String(self.cmd_id.clone()));
76 if let Some(s) = signal {
77 p.insert("signal".into(), Value::from(s));
78 }
79 self.channel.call("cmd.kill", Value::Object(p)).await.map(|_| ())
80 }
81
82 pub async fn wait(self) -> Result<CommandResult, SolariError> {
84 self.wait_deadline(None).await
85 }
86
87 pub(crate) async fn wait_deadline(
90 mut self,
91 deadline: Option<tokio::time::Instant>,
92 ) -> Result<CommandResult, SolariError> {
93 let mut stdout = String::new();
94 let mut stderr = String::new();
95 loop {
96 let next = match deadline {
97 Some(d) => match tokio::time::timeout_at(d, self.rx.recv()).await {
98 Ok(frame) => frame,
99 Err(_) => {
100 let _ = self.kill(Some(9)).await;
101 self.channel.unregister_command(&self.cmd_id);
102 return Ok(CommandResult { exit_code: -1, stdout, stderr });
103 }
104 },
105 None => self.rx.recv().await,
106 };
107 let Some(frame) = next else { break };
108 match frame.typ.as_str() {
109 "cmd.data" => {
110 let text = frame
111 .base64
112 .as_ref()
113 .map(|b| decode_text(b))
114 .unwrap_or_default();
115 let is_err = frame.stream.as_deref() == Some("stderr");
116 if is_err {
117 stderr.push_str(&text);
118 if let Some(cb) = &self.on_stderr {
119 cb(&text);
120 }
121 } else {
122 stdout.push_str(&text);
123 if let Some(cb) = &self.on_stdout {
124 cb(&text);
125 }
126 }
127 }
128 "cmd.exit" => {
129 let exit_code = frame.exit_code.unwrap_or(0);
130 self.channel.unregister_command(&self.cmd_id);
131 return Ok(CommandResult { exit_code, stdout, stderr });
132 }
133 _ => {}
134 }
135 }
136 Err(SolariError::connection("Control channel closed"))
138 }
139}
140
141fn decode_text(b64: &str) -> String {
142 match base64::engine::general_purpose::STANDARD.decode(b64) {
143 Ok(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
144 Err(_) => String::new(),
145 }
146}
147
148fn cmd_start_params(cmd: &str, opts: &RunOptions) -> Value {
151 let mut p = Map::new();
152 p.insert("cmd".into(), Value::String(cmd.to_string()));
153 if let Some(args) = &opts.args {
154 p.insert("args".into(), Value::from(args.clone()));
155 }
156 if let Some(cwd) = &opts.cwd {
157 p.insert("cwd".into(), Value::String(cwd.clone()));
158 }
159 if let Some(env) = &opts.env {
160 p.insert("env".into(), serde_json::to_value(env).unwrap());
161 }
162 if let Some(user) = &opts.user {
163 p.insert("user".into(), Value::String(user.clone()));
164 }
165 Value::Object(p)
166}
167
168pub(crate) async fn start_command(
170 sb: &Sandbox,
171 cmd: &str,
172 opts: &RunOptions,
173) -> Result<CommandHandle, SolariError> {
174 sb.channel().connect().await?;
175 let result = sb
176 .channel()
177 .call("cmd.start", cmd_start_params(cmd, opts))
178 .await?;
179 let cmd_id = result
180 .get("cmdId")
181 .and_then(Value::as_str)
182 .ok_or_else(|| SolariError::Other("cmd.start reply missing cmdId".into()))?
183 .to_string();
184 let rx = sb.channel().register_command(&cmd_id);
185 Ok(CommandHandle {
186 cmd_id,
187 channel: sb.channel_arc(),
188 rx,
189 on_stdout: opts.on_stdout.clone(),
190 on_stderr: opts.on_stderr.clone(),
191 })
192}
193
194pub(crate) async fn run_command(
196 sb: &Sandbox,
197 cmd: &str,
198 opts: &RunOptions,
199) -> Result<CommandResult, SolariError> {
200 if !opts.background
204 && opts.on_stdout.is_none()
205 && opts.on_stderr.is_none()
206 && opts.env.is_none()
207 && opts.user.is_none()
208 && !sb.channel().connected()
209 {
210 match exec_oneshot(sb, cmd, opts).await {
211 Ok(r) => return Ok(r),
212 Err(err) => {
215 if !is_route_unavailable(&err) {
216 return Err(err);
217 }
218 }
219 }
220 }
221
222 let handle = start_command(sb, cmd, opts).await?;
223 if opts.background {
224 return Ok(CommandResult {
225 exit_code: 0,
226 stdout: String::new(),
227 stderr: String::new(),
228 });
229 }
230 let deadline = opts
232 .timeout_ms
233 .filter(|t| *t > 0)
234 .map(|t| tokio::time::Instant::now() + std::time::Duration::from_millis(t));
235 handle.wait_deadline(deadline).await
236}
237
238async fn exec_oneshot(
240 sb: &Sandbox,
241 cmd: &str,
242 opts: &RunOptions,
243) -> Result<CommandResult, SolariError> {
244 let mut body = Map::new();
245 body.insert("cmd".into(), Value::String(cmd.to_string()));
246 if let Some(args) = &opts.args {
247 body.insert("args".into(), Value::from(args.clone()));
248 }
249 if let Some(cwd) = &opts.cwd {
250 body.insert("cwd".into(), Value::String(cwd.clone()));
251 }
252 if let Some(t) = opts.timeout_ms {
253 body.insert("timeoutMs".into(), Value::from(t));
254 }
255 let path = format!("/sandboxes/{}/exec", encode_uri_component(sb.id()));
256 sb.http()
257 .request::<CommandResult>("POST", &path, Some(Value::Object(body)), None)
258 .await
259}
260
261pub struct Commands<'a> {
263 pub(crate) sb: &'a Sandbox,
264}
265
266impl<'a> Commands<'a> {
267 pub async fn run(&self, cmd: &str, opts: RunOptions) -> Result<CommandResult, SolariError> {
269 run_command(self.sb, cmd, &opts).await
270 }
271
272 pub async fn start(&self, cmd: &str, opts: RunOptions) -> Result<CommandHandle, SolariError> {
274 start_command(self.sb, cmd, &opts).await
275 }
276}