aa-cli 0.0.1-alpha.9

aasm — command-line tool for Agent Assembly
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

//! `aasm sandbox` subcommand — execute WebAssembly tools inside the
//! Agent Assembly tool-execution sandbox (highlight ④ of the product
//! spec). Surfaces [`aa_sandbox::SandboxRuntime`] to OSS users so the
//! filesystem + CPU + memory + wall-clock budgets enforced by
//! `aa-sandbox` are reachable from the `aasm` CLI without going
//! through the cloud `aa-api` `/dispatch_tool` HTTP route.
//!
//! Issue: AAASM-2340.

use std::path::PathBuf;
use std::process::ExitCode;

use aa_sandbox::policy::{SandboxConfig, SandboxLimits};
use aa_sandbox::runtime::SandboxRuntime;

/// Arguments for `aasm sandbox`.
#[derive(Debug, clap::Args)]
pub struct SandboxArgs {
    #[command(subcommand)]
    pub subcommand: SandboxSubcommand,
}

/// Subcommands available under `aasm sandbox`.
#[derive(Debug, clap::Subcommand)]
pub enum SandboxSubcommand {
    /// Run a WebAssembly module inside a fresh sandbox and report the outcome.
    Run(RunArgs),
    /// Show the default sandbox runtime limits.
    Info,
}

/// Arguments for `aasm sandbox run`.
#[derive(Debug, clap::Args)]
pub struct RunArgs {
    /// Path to a `.wasm` module to execute under WASI preview 1.
    pub wasm: PathBuf,

    /// Wasmtime instruction-fuel budget. Defaults to the safe-by-default
    /// 10M units; pass a larger value for long-running tools.
    #[arg(long)]
    pub fuel: Option<u64>,

    /// Maximum linear-memory pages (1 page = 64 KiB). Defaults to 16 (1 MiB).
    #[arg(long)]
    pub memory_pages: Option<u32>,

    /// Wall-clock deadline in milliseconds. Defaults to 5000 (5s).
    #[arg(long)]
    pub wall_clock_ms: Option<u64>,
}

/// Dispatch the `sandbox` subcommand.
pub fn dispatch(args: SandboxArgs) -> ExitCode {
    match args.subcommand {
        SandboxSubcommand::Run(run) => run_wasm(run),
        SandboxSubcommand::Info => print_info(),
    }
}

fn run_wasm(args: RunArgs) -> ExitCode {
    let bytes = match std::fs::read(&args.wasm) {
        Ok(b) => b,
        Err(e) => {
            eprintln!("error: failed to read {}: {e}", args.wasm.display());
            return ExitCode::FAILURE;
        }
    };

    let limits = SandboxLimits {
        fuel: args.fuel.unwrap_or(SandboxLimits::default().fuel),
        memory_pages: args.memory_pages.unwrap_or(SandboxLimits::default().memory_pages),
        wall_clock_ms: args.wall_clock_ms.unwrap_or(SandboxLimits::default().wall_clock_ms),
    };
    let config = SandboxConfig {
        preopened_dirs: Vec::new(),
        limits,
    };

    let runtime = match SandboxRuntime::new(config) {
        Ok(r) => r,
        Err(e) => {
            eprintln!("error: failed to build sandbox runtime: {e}");
            return ExitCode::FAILURE;
        }
    };

    match runtime.run_tool(&bytes, &[]) {
        Ok(output) => {
            println!("sandbox exited cleanly (exit_code={})", output.exit_code);
            ExitCode::SUCCESS
        }
        Err(e) => {
            eprintln!("sandbox refused or trapped the module: {e}");
            ExitCode::FAILURE
        }
    }
}

fn print_info() -> ExitCode {
    let limits = SandboxLimits::default();
    println!("aasm sandbox — WASI preview 1 tool-execution sandbox");
    println!("  fuel (instructions):      {}", limits.fuel);
    println!(
        "  memory ceiling:           {} pages ({} KiB)",
        limits.memory_pages,
        (limits.memory_pages as usize) * 64
    );
    println!("  wall-clock deadline (ms): {}", limits.wall_clock_ms);
    println!("  preopened dirs:           (none — fully sealed FS)");
    ExitCode::SUCCESS
}