tinysandbox 0.1.0

Ultra-minimal Linux-like sandbox for AI agents: shell, coreutils, VFS, and a Wasmtime-hosted JS runtime in one crate
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

use std::collections::{BTreeMap, BTreeSet};
use std::future::Future;
use std::io;
use std::pin::Pin;
use std::sync::{Arc, Mutex, PoisonError};
use std::task::{Context, Poll};
use std::time::Duration;

use tokio::io::{AsyncRead, AsyncWrite};

use super::fs::Fs;

pub type CommandFuture = Pin<Box<dyn Future<Output = CommandResult> + Send>>;
pub type BoxAsyncRead = Pin<Box<dyn AsyncRead + Send>>;
pub type BoxAsyncWrite = Pin<Box<dyn AsyncWrite + Send>>;

pub trait Command: Send + Sync {
    fn run(&self, ctx: CommandContext) -> CommandFuture;
}

impl<F, Fut> Command for F
where
    F: Fn(CommandContext) -> Fut + Send + Sync,
    Fut: Future<Output = CommandResult> + Send + 'static,
{
    fn run(&self, ctx: CommandContext) -> CommandFuture {
        Box::pin(self(ctx))
    }
}

pub struct CommandContext {
    pub args: Vec<String>,
    pub env: BTreeMap<String, String>,
    pub cwd: String,
    pub stdin: BoxAsyncRead,
    pub stdout: BoxAsyncWrite,
    pub stderr: BoxAsyncWrite,
    pub fs: Fs,
    pub limits: Limits,
    pub commands: Arc<BTreeSet<String>>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CommandResult {
    pub exit_code: i32,
    pub peak_wasm_memory_bytes: Option<usize>,
}

impl CommandResult {
    pub const fn new(exit_code: i32) -> Self {
        Self {
            exit_code,
            peak_wasm_memory_bytes: None,
        }
    }

    pub const fn with_peak_wasm_memory(mut self, bytes: usize) -> Self {
        self.peak_wasm_memory_bytes = Some(bytes);
        self
    }

    pub const fn success() -> Self {
        Self::new(0)
    }

    pub const fn failure() -> Self {
        Self::new(1)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Limits {
    pub wall_time: Duration,
    pub stdout_bytes: usize,
    pub stderr_bytes: usize,
    pub max_commands: usize,
    pub sort_input_bytes: usize,
    pub wasm_memory_bytes: usize,
}

impl Default for Limits {
    fn default() -> Self {
        Self {
            wall_time: Duration::from_secs(30),
            stdout_bytes: 1024 * 1024,
            stderr_bytes: 1024 * 1024,
            max_commands: 1024,
            sort_input_bytes: 8 * 1024 * 1024,
            wasm_memory_bytes: 64 * 1024 * 1024,
        }
    }
}

#[derive(Clone)]
pub(crate) struct SharedWriter {
    inner: Arc<Mutex<Vec<u8>>>,
}

impl SharedWriter {
    pub(crate) fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(Vec::new())),
        }
    }

    pub(crate) fn boxed(&self) -> BoxAsyncWrite {
        Box::pin(self.clone())
    }

    pub(crate) fn bytes(&self) -> Vec<u8> {
        self.inner
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .clone()
    }
}

impl AsyncWrite for SharedWriter {
    fn poll_write(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        self.inner
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .extend_from_slice(buf);
        Poll::Ready(Ok(buf.len()))
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }
}