Skip to main content

solari/
commands.rs

1//! `commands` namespace: run / start / stdin / kill over the control WS, plus
2//! the one-shot REST `/exec` fast path. Mirrors the `commands` surface + the
3//! `runCommand`/`startCommand` internals of `handle.ts`.
4
5use 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
17/// A streamed-output callback (`stdout`/`stderr`). `Fn(&str)`.
18pub type OutputCallback = Arc<dyn Fn(&str) + Send + Sync>;
19
20/// Options for `commands.run` / `commands.start`.
21#[derive(Default, Clone)]
22pub struct RunOptions {
23    /// argv tail (the guest runs `cmd` with these, NOT via a shell).
24    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    /// Run detached: returns immediately, output still streams via callbacks.
30    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
53/// A started command handle (`commands.start`).
54pub 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    /// Write bytes to the command's stdin.
64    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    /// Send a signal (default SIGTERM guest-side).
73    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    /// Wait for the command to exit, accumulating stdout/stderr.
83    pub async fn wait(self) -> Result<CommandResult, SolariError> {
84        self.wait_deadline(None).await
85    }
86
87    /// `wait` with an optional deadline. When it elapses the command is killed (SIGKILL) and the
88    /// result is exit -1 with the output so far, the same result the guest's own `/exec` timeout gives.
89    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        // Channel dropped before `cmd.exit`.
137        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
148/// Build a JSON params object, inserting only present (`Some`) fields — mirrors
149/// how JSON.stringify omits `undefined`.
150fn 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
168/// `commands.start`: start a process and return a handle immediately.
169pub(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
194/// `commands.run`: run a command to completion.
195pub(crate) async fn run_command(
196    sb: &Sandbox,
197    cmd: &str,
198    opts: &RunOptions,
199) -> Result<CommandResult, SolariError> {
200    // One-shot HTTP fast path: a plain run-to-completion (no streaming
201    // callbacks, not backgrounded, no per-command env/user) over the warm REST
202    // connection, ONLY when the control channel isn't already open.
203    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            // Fall back to the WS path ONLY when `/exec` is unavailable. Any
213            // other failure propagates (never silently double-execute).
214            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    // `timeout_ms` is a deadline for the command here too (it used to be honoured only by the REST path).
231    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
238/// POST `/sandboxes/:id/exec` — the one-shot fast path.
239async 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
261/// The ergonomic `commands` accessor returned by `Sandbox::commands()`.
262pub struct Commands<'a> {
263    pub(crate) sb: &'a Sandbox,
264}
265
266impl<'a> Commands<'a> {
267    /// Run a command to completion.
268    pub async fn run(&self, cmd: &str, opts: RunOptions) -> Result<CommandResult, SolariError> {
269        run_command(self.sb, cmd, &opts).await
270    }
271
272    /// Start a command and return a handle immediately.
273    pub async fn start(&self, cmd: &str, opts: RunOptions) -> Result<CommandHandle, SolariError> {
274        start_command(self.sb, cmd, &opts).await
275    }
276}