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/// How long a cancelled command gets to exit after SIGTERM before it is SIGKILLed.
54const CANCEL_KILL_GRACE: std::time::Duration = std::time::Duration::from_secs(2);
55
56/// A cloneable handle on a running command's input side: write or close stdin and send signals while another task is
57/// blocked in [`CommandHandle::wait`] (which consumes the handle). Get one with [`CommandHandle::io`].
58#[derive(Clone)]
59pub struct CommandIo {
60    cmd_id: String,
61    channel: Arc<ControlChannel>,
62}
63
64impl CommandIo {
65    /// The command id.
66    pub fn cmd_id(&self) -> &str {
67        &self.cmd_id
68    }
69
70    /// Write bytes to the command's stdin.
71    pub async fn stdin(&self, data: impl AsRef<[u8]>) -> Result<(), SolariError> {
72        let b64 = base64::engine::general_purpose::STANDARD.encode(data.as_ref());
73        let mut p = Map::new();
74        p.insert("cmdId".into(), Value::String(self.cmd_id.clone()));
75        p.insert("base64".into(), Value::String(b64));
76        self.channel.call("cmd.stdin", Value::Object(p)).await.map(|_| ())
77    }
78
79    /// Close the command's stdin (the guest writes nothing, then closes it).
80    pub async fn close_stdin(&self) -> Result<(), SolariError> {
81        let mut p = Map::new();
82        p.insert("cmdId".into(), Value::String(self.cmd_id.clone()));
83        p.insert("base64".into(), Value::String(String::new()));
84        p.insert("eof".into(), Value::Bool(true));
85        self.channel.call("cmd.stdin", Value::Object(p)).await.map(|_| ())
86    }
87
88    /// Send a signal (default SIGTERM guest-side).
89    pub async fn kill(&self, signal: Option<i64>) -> Result<(), SolariError> {
90        let mut p = Map::new();
91        p.insert("cmdId".into(), Value::String(self.cmd_id.clone()));
92        if let Some(s) = signal {
93            p.insert("signal".into(), Value::from(s));
94        }
95        self.channel.call("cmd.kill", Value::Object(p)).await.map(|_| ())
96    }
97}
98
99/// Stops the guest process when a `commands.run` future is dropped before the command finished: SIGTERM and a closed stdin
100/// at once, SIGKILL after [`CANCEL_KILL_GRACE`]. Disarmed once the run returns. Every call is best effort (a command that
101/// already finished answers with an error).
102struct CancelGuard {
103    io: Option<CommandIo>,
104}
105
106impl Drop for CancelGuard {
107    fn drop(&mut self) {
108        let Some(io) = self.io.take() else { return };
109        let Ok(rt) = tokio::runtime::Handle::try_current() else { return };
110        rt.spawn(async move {
111            let _ = io.kill(Some(15)).await;
112            let _ = io.close_stdin().await;
113            tokio::time::sleep(CANCEL_KILL_GRACE).await;
114            let _ = io.kill(Some(9)).await;
115        });
116    }
117}
118
119/// A started command handle (`commands.start`).
120pub struct CommandHandle {
121    pub cmd_id: String,
122    channel: Arc<ControlChannel>,
123    rx: mpsc::UnboundedReceiver<AsyncFrame>,
124    on_stdout: Option<OutputCallback>,
125    on_stderr: Option<OutputCallback>,
126}
127
128impl CommandHandle {
129    /// A cloneable handle on this command's input side. [`wait`](Self::wait) consumes the handle, so take this first to
130    /// keep writing stdin or signalling while another task waits.
131    pub fn io(&self) -> CommandIo {
132        CommandIo { cmd_id: self.cmd_id.clone(), channel: self.channel.clone() }
133    }
134
135    /// Write bytes to the command's stdin.
136    pub async fn stdin(&self, data: impl AsRef<[u8]>) -> Result<(), SolariError> {
137        let b64 = base64::engine::general_purpose::STANDARD.encode(data.as_ref());
138        let mut p = Map::new();
139        p.insert("cmdId".into(), Value::String(self.cmd_id.clone()));
140        p.insert("base64".into(), Value::String(b64));
141        self.channel.call("cmd.stdin", Value::Object(p)).await.map(|_| ())
142    }
143
144    /// Send a signal (default SIGTERM guest-side).
145    pub async fn kill(&self, signal: Option<i64>) -> Result<(), SolariError> {
146        let mut p = Map::new();
147        p.insert("cmdId".into(), Value::String(self.cmd_id.clone()));
148        if let Some(s) = signal {
149            p.insert("signal".into(), Value::from(s));
150        }
151        self.channel.call("cmd.kill", Value::Object(p)).await.map(|_| ())
152    }
153
154    /// Wait for the command to exit, accumulating stdout/stderr.
155    pub async fn wait(self) -> Result<CommandResult, SolariError> {
156        self.wait_deadline(None).await
157    }
158
159    /// `wait` with an optional deadline. When it elapses the command is killed (SIGKILL) and the
160    /// result is exit -1 with the output so far, the same result the guest's own `/exec` timeout gives.
161    pub(crate) async fn wait_deadline(
162        mut self,
163        deadline: Option<tokio::time::Instant>,
164    ) -> Result<CommandResult, SolariError> {
165        let mut stdout = String::new();
166        let mut stderr = String::new();
167        let mut out_dec = Utf8Stream::default();
168        let mut err_dec = Utf8Stream::default();
169        loop {
170            let next = match deadline {
171                Some(d) => match tokio::time::timeout_at(d, self.rx.recv()).await {
172                    Ok(frame) => frame,
173                    Err(_) => {
174                        let _ = self.kill(Some(9)).await;
175                        self.channel.unregister_command(&self.cmd_id);
176                        return Ok(CommandResult { exit_code: -1, stdout, stderr, timed_out: true });
177                    }
178                },
179                None => self.rx.recv().await,
180            };
181            let Some(frame) = next else { break };
182            match frame.typ.as_str() {
183                "cmd.data" => {
184                    let is_err = frame.stream.as_deref() == Some("stderr");
185                    // One streaming decoder per stream: a multi-byte character split across two frames must come out
186                    // whole, not as U+FFFD at each end.
187                    let bytes = frame.base64.as_ref().map(|b| decode_bytes(b)).unwrap_or_default();
188                    let text = if is_err { err_dec.push(&bytes) } else { out_dec.push(&bytes) };
189                    if is_err {
190                        stderr.push_str(&text);
191                        if let Some(cb) = &self.on_stderr {
192                            cb(&text);
193                        }
194                    } else {
195                        stdout.push_str(&text);
196                        if let Some(cb) = &self.on_stdout {
197                            cb(&text);
198                        }
199                    }
200                }
201                "cmd.exit" => {
202                    let exit_code = frame.exit_code.unwrap_or(0);
203                    self.channel.unregister_command(&self.cmd_id);
204                    // Bytes held back at the end of a stream (a truncated character) are flushed, not dropped.
205                    let rest = out_dec.finish();
206                    if !rest.is_empty() {
207                        stdout.push_str(&rest);
208                        if let Some(cb) = &self.on_stdout {
209                            cb(&rest);
210                        }
211                    }
212                    let rest = err_dec.finish();
213                    if !rest.is_empty() {
214                        stderr.push_str(&rest);
215                        if let Some(cb) = &self.on_stderr {
216                            cb(&rest);
217                        }
218                    }
219                    return Ok(CommandResult { exit_code, stdout, stderr, timed_out: false });
220                }
221                _ => {}
222            }
223        }
224        // Channel dropped before `cmd.exit`.
225        Err(SolariError::connection("Control channel closed"))
226    }
227}
228
229fn decode_bytes(b64: &str) -> Vec<u8> {
230    base64::engine::general_purpose::STANDARD.decode(b64).unwrap_or_default()
231}
232
233/// Streaming UTF-8 decoder: holds back an incomplete character at the end of a chunk until the next chunk completes it.
234/// Invalid bytes become U+FFFD, as `String::from_utf8_lossy` does.
235#[derive(Default)]
236pub(crate) struct Utf8Stream {
237    pending: Vec<u8>,
238}
239
240impl Utf8Stream {
241    pub(crate) fn push(&mut self, bytes: &[u8]) -> String {
242        self.pending.extend_from_slice(bytes);
243        let mut out = String::new();
244        loop {
245            match std::str::from_utf8(&self.pending) {
246                Ok(s) => {
247                    out.push_str(s);
248                    self.pending.clear();
249                    break;
250                }
251                Err(e) => {
252                    let valid = e.valid_up_to();
253                    out.push_str(std::str::from_utf8(&self.pending[..valid]).unwrap_or(""));
254                    match e.error_len() {
255                        Some(n) => {
256                            out.push('\u{FFFD}');
257                            self.pending.drain(..valid + n);
258                        }
259                        None => {
260                            // An incomplete character at the end: wait for the next chunk.
261                            self.pending.drain(..valid);
262                            break;
263                        }
264                    }
265                }
266            }
267        }
268        out
269    }
270
271    /// End of stream: a character still incomplete is reported once as U+FFFD.
272    pub(crate) fn finish(&mut self) -> String {
273        if self.pending.is_empty() {
274            String::new()
275        } else {
276            self.pending.clear();
277            "\u{FFFD}".to_string()
278        }
279    }
280}
281
282#[cfg(test)]
283mod utf8_tests {
284    use super::Utf8Stream;
285
286    #[test]
287    fn characters_split_across_chunks_come_out_whole() {
288        let text = "漢字🙂é".repeat(50);
289        let mut dec = Utf8Stream::default();
290        let mut out = String::new();
291        for b in text.as_bytes() {
292            out.push_str(&dec.push(&[*b]));
293        }
294        out.push_str(&dec.finish());
295        assert_eq!(out, text);
296        assert!(!out.contains('\u{FFFD}'));
297    }
298
299    #[test]
300    fn a_truncated_character_at_the_end_is_reported_once_and_invalid_bytes_are_replaced() {
301        let mut dec = Utf8Stream::default();
302        assert_eq!(dec.push(b"ok\xe6"), "ok");
303        assert_eq!(dec.finish(), "\u{FFFD}");
304        let mut dec = Utf8Stream::default();
305        assert_eq!(dec.push(b"a\xffb"), "a\u{FFFD}b");
306        assert_eq!(dec.finish(), "");
307    }
308}
309
310/// Build a JSON params object, inserting only present (`Some`) fields — mirrors
311/// how JSON.stringify omits `undefined`.
312fn cmd_start_params(cmd: &str, opts: &RunOptions) -> Value {
313    let mut p = Map::new();
314    p.insert("cmd".into(), Value::String(cmd.to_string()));
315    if let Some(args) = &opts.args {
316        p.insert("args".into(), Value::from(args.clone()));
317    }
318    if let Some(cwd) = &opts.cwd {
319        p.insert("cwd".into(), Value::String(cwd.clone()));
320    }
321    if let Some(env) = &opts.env {
322        p.insert("env".into(), serde_json::to_value(env).unwrap());
323    }
324    if let Some(user) = &opts.user {
325        p.insert("user".into(), Value::String(user.clone()));
326    }
327    Value::Object(p)
328}
329
330/// `commands.start`: start a process and return a handle immediately.
331pub(crate) async fn start_command(
332    sb: &Sandbox,
333    cmd: &str,
334    opts: &RunOptions,
335) -> Result<CommandHandle, SolariError> {
336    sb.channel().connect().await?;
337    let result = sb
338        .channel()
339        .call("cmd.start", cmd_start_params(cmd, opts))
340        .await?;
341    let cmd_id = result
342        .get("cmdId")
343        .and_then(Value::as_str)
344        .ok_or_else(|| SolariError::Other("cmd.start reply missing cmdId".into()))?
345        .to_string();
346    let rx = sb.channel().register_command(&cmd_id);
347    Ok(CommandHandle {
348        cmd_id,
349        channel: sb.channel_arc(),
350        rx,
351        on_stdout: opts.on_stdout.clone(),
352        on_stderr: opts.on_stderr.clone(),
353    })
354}
355
356/// `commands.run`: run a command to completion.
357pub(crate) async fn run_command(
358    sb: &Sandbox,
359    cmd: &str,
360    opts: &RunOptions,
361) -> Result<CommandResult, SolariError> {
362    // One-shot HTTP fast path: a plain run-to-completion (no streaming
363    // callbacks, not backgrounded, no per-command env/user) over the warm REST
364    // connection, ONLY when the control channel isn't already open.
365    if !opts.background
366        && opts.on_stdout.is_none()
367        && opts.on_stderr.is_none()
368        && opts.env.is_none()
369        && opts.user.is_none()
370        && !sb.channel().connected()
371    {
372        match exec_oneshot(sb, cmd, opts).await {
373            Ok(r) => return Ok(r),
374            // Fall back to the WS path ONLY when `/exec` is unavailable. Any
375            // other failure propagates (never silently double-execute).
376            Err(err) => {
377                if !is_route_unavailable(&err) {
378                    return Err(err);
379                }
380            }
381        }
382    }
383
384    let handle = start_command(sb, cmd, opts).await?;
385    if opts.background {
386        // Nobody awaits a background command, so something has to drive its output to the callbacks: they were never called.
387        if opts.on_stdout.is_some() || opts.on_stderr.is_some() {
388            tokio::spawn(async move {
389                let _ = handle.wait().await;
390            });
391        }
392        return Ok(CommandResult {
393            exit_code: 0,
394            stdout: String::new(),
395            stderr: String::new(),
396            timed_out: false,
397        });
398    }
399    // `timeout_ms` is a deadline for the command here too (it used to be honoured only by the REST path).
400    let deadline = opts
401        .timeout_ms
402        .filter(|t| *t > 0)
403        .map(|t| tokio::time::Instant::now() + std::time::Duration::from_millis(t));
404    let mut guard = CancelGuard { io: Some(handle.io()) };
405    let result = handle.wait_deadline(deadline).await;
406    guard.io = None; // finished (or failed) on its own: nothing to stop
407    result
408}
409
410/// POST `/sandboxes/:id/exec` — the one-shot fast path.
411async fn exec_oneshot(
412    sb: &Sandbox,
413    cmd: &str,
414    opts: &RunOptions,
415) -> Result<CommandResult, SolariError> {
416    let mut body = Map::new();
417    body.insert("cmd".into(), Value::String(cmd.to_string()));
418    if let Some(args) = &opts.args {
419        body.insert("args".into(), Value::from(args.clone()));
420    }
421    if let Some(cwd) = &opts.cwd {
422        body.insert("cwd".into(), Value::String(cwd.clone()));
423    }
424    if let Some(t) = opts.timeout_ms {
425        body.insert("timeoutMs".into(), Value::from(t));
426    }
427    let path = format!("/sandboxes/{}/exec", encode_uri_component(sb.id()));
428    sb.http()
429        .request::<CommandResult>("POST", &path, Some(Value::Object(body)), None)
430        .await
431}
432
433/// The ergonomic `commands` accessor returned by `Sandbox::commands()`.
434pub struct Commands<'a> {
435    pub(crate) sb: &'a Sandbox,
436}
437
438impl<'a> Commands<'a> {
439    /// Run a command to completion.
440    pub async fn run(&self, cmd: &str, opts: RunOptions) -> Result<CommandResult, SolariError> {
441        run_command(self.sb, cmd, &opts).await
442    }
443
444    /// Start a command and return a handle immediately.
445    pub async fn start(&self, cmd: &str, opts: RunOptions) -> Result<CommandHandle, SolariError> {
446        start_command(self.sb, cmd, &opts).await
447    }
448}
solari/commands.rs

solari/
commands.rs
