@@ -1,7 +1,3 @@
1−mod session_manager;
2−
3−pub use session_manager::{SessionError, SessionManager, SessionRecord};
4−
51use std::{
62 env,
73 ffi::OsString,
@@ -119,41 +115,26 @@ impl DirectoryShare {
119115 }
120116}
121117
122−fn main() -> Result<(), Box<dyn std::error::Error>> {
118+pub fn run_cli() -> Result<(), Box<dyn std::error::Error>> {
123119 let args = parse_cli()?;
124120
125− if args.version {
126− println!("Vibe");
127− println!("https://github.com/lynaghk/vibe/");
128− println!("Git SHA: {}", env!("GIT_SHA"));
129− std::process::exit(0);
121+ if args.version() {
122+ print_version();
123+ return Ok(());
130124 }
131125
132− if args.help {
133− println!(
134− "Vibe is a quick way to spin up a Linux virtual machine on Mac to sandbox LLM agents.
135−
136−vibe [OPTIONS] [disk-image.raw]
137−
138−Options
139−
140− --help Print this help message.
141− --version Print the version (commit SHA).
142− --no-default-mounts Disable all default mounts.
143− --mount host-path:guest-path[:read-only | :read-write] Mount `host-path` inside VM at `guest-path`.
144− Defaults to read-write.
145− Errors if host-path does not exist.
146− --cpus <count> Number of virtual CPUs (default {DEFAULT_CPU_COUNT}).
147− --ram <megabytes> RAM size in megabytes (default {DEFAULT_RAM_MB}).
148− --script <path/to/script.sh> Run script in VM.
149− --send <some-command> Type `some-command` followed by newline into the VM.
150− --expect <string> [timeout-seconds] Wait for `string` to appear in console output before executing next `--script` or `--send`.
151− If `string` does not appear within timeout (default 30 seconds), shutdown VM with error.
152−"
153− );
154− std::process::exit(0);
126+ if args.help() {
127+ print_help();
128+ return Ok(());
155129 }
156130
131+ run_with_args(args, spawn_vm_io)
132+}
133+
134+pub fn run_with_args<F>(args: CliArgs, io_handler: F) -> Result<(), Box<dyn std::error::Error>>
135+where
136+ F: FnOnce(Arc<OutputMonitor>, OwnedFd, OwnedFd) -> IoContext,
137+{
157138 ensure_signed();
158139
159140 let project_root = env::current_dir()?;
@@ -239,6 +220,11 @@ Options
239220 ),
240221 DirectoryShare::new(home.join(".codex"), "/root/.codex".into(), false),
241222 DirectoryShare::new(home.join(".claude"), "/root/.claude".into(), false),
223+ DirectoryShare::new(
224+ "/Users/zhangjie/Documents/Code/CompletePrograms/vibebox/.ssh".into(),
225+ "/root/.ssh".into(),
226+ true,
227+ ),
242228 ]
243229 .into_iter()
244230 .flatten()
@@ -258,16 +244,47 @@ Options
258244 // Any user-provided login actions must come after our system ones
259245 login_actions.extend(args.login_actions);
260246
261− run_vm(
247+ run_vm_with_io(
262248 &disk_path,
263249 &login_actions,
264250 &directory_shares[..],
265251 args.cpu_count,
266252 args.ram_bytes,
253+ io_handler,
267254 )
268255}
269256
270−struct CliArgs {
257+pub fn print_help() {
258+ println!(
259+ "Vibe is a quick way to spin up a Linux virtual machine on Mac to sandbox LLM agents.
260+
261+vibe [OPTIONS] [disk-image.raw]
262+
263+Options
264+
265+ --help Print this help message.
266+ --version Print the version (commit SHA).
267+ --no-default-mounts Disable all default mounts.
268+ --mount host-path:guest-path[:read-only | :read-write] Mount `host-path` inside VM at `guest-path`.
269+ Defaults to read-write.
270+ Errors if host-path does not exist.
271+ --cpus <count> Number of virtual CPUs (default {DEFAULT_CPU_COUNT}).
272+ --ram <megabytes> RAM size in megabytes (default {DEFAULT_RAM_MB}).
273+ --script <path/to/script.sh> Run script in VM.
274+ --send <some-command> Type `some-command` followed by newline into the VM.
275+ --expect <string> [timeout-seconds] Wait for `string` to appear in console output before executing next `--script` or `--send`.
276+ If `string` does not appear within timeout (default 30 seconds), shutdown VM with error.
277+"
278+ );
279+}
280+
281+pub fn print_version() {
282+ println!("Vibe");
283+ println!("https://github.com/lynaghk/vibe/");
284+ println!("Git SHA: {}", env!("GIT_SHA"));
285+}
286+
287+pub struct CliArgs {
271288 disk: Option<PathBuf>,
272289 version: bool,
273290 help: bool,
@@ -278,14 +295,43 @@ struct CliArgs {
278295 ram_bytes: u64,
279296}
280297
281−fn parse_cli() -> Result<CliArgs, Box<dyn std::error::Error>> {
298+impl CliArgs {
299+ pub fn version(&self) -> bool {
300+ self.version
301+ }
302+
303+ pub fn help(&self) -> bool {
304+ self.help
305+ }
306+
307+ pub fn cpu_count(&self) -> usize {
308+ self.cpu_count
309+ }
310+
311+ pub fn ram_bytes(&self) -> u64 {
312+ self.ram_bytes
313+ }
314+
315+ pub fn ram_mb(&self) -> u64 {
316+ self.ram_bytes / BYTES_PER_MB
317+ }
318+}
319+
320+pub fn parse_cli() -> Result<CliArgs, Box<dyn std::error::Error>> {
321+ parse_cli_from(env::args_os())
322+}
323+
324+pub fn parse_cli_from<I>(args: I) -> Result<CliArgs, Box<dyn std::error::Error>>
325+where
326+ I: IntoIterator<Item = OsString>,
327+{
282328 fn os_to_string(value: OsString, flag: &str) -> Result<String, Box<dyn std::error::Error>> {
283329 value
284330 .into_string()
285331 .map_err(|_| format!("{flag} expects valid UTF-8").into())
286332 }
287333
288− let mut parser = lexopt::Parser::from_env();
334+ let mut parser = lexopt::Parser::from_iter(args);
289335 let mut disk = None;
290336 let mut version = false;
291337 let mut help = false;
@@ -621,11 +667,15 @@ pub fn create_pipe() -> (OwnedFd, OwnedFd) {
621667 (read_stream.into(), write_stream.into())
622668}
623669
624−pub fn spawn_vm_io(
670+pub fn spawn_vm_io_with_line_handler<F>(
625671 output_monitor: Arc<OutputMonitor>,
626672 vm_output_fd: OwnedFd,
627673 vm_input_fd: OwnedFd,
628−) -> IoContext {
674+ mut on_line: F,
675+) -> IoContext
676+where
677+ F: FnMut(&str) -> bool + ::std::marker::Send + 'static,
678+{
629679 let (input_tx, input_rx): (Sender<VmInput>, Receiver<VmInput>) = mpsc::channel();
630680
631681 // raw_guard is set when we've put the user's terminal into raw mode because we've attached stdin/stdout to the VM.
@@ -679,15 +729,41 @@ pub fn spawn_vm_io(
679729
680730 move || {
681731 let mut buf = [0u8; 1024];
732+ let mut pending_command: Vec<u8> = Vec::new();
733+ let mut command_mode = false;
682734 loop {
683735 match poll_with_wakeup(libc::STDIN_FILENO, wakeup_read.as_raw_fd(), &mut buf) {
684736 PollResult::Shutdown | PollResult::Error => break,
685737 PollResult::Spurious => continue,
686738 PollResult::Ready(bytes) => {
739+ let mut send_buf: Vec<u8> = Vec::new();
740+ for &b in bytes {
741+ if pending_command.is_empty() && !command_mode && b == b':' {
742+ command_mode = true;
743+ }
744+
745+ if command_mode {
746+ pending_command.push(b);
747+ } else {
748+ send_buf.push(b);
749+ }
750+
751+ if b == b'\n' && command_mode {
752+ let line = String::from_utf8_lossy(&pending_command);
753+ let trimmed = line.trim_end_matches(&['\r', '\n'][..]);
754+ let consumed = on_line(trimmed);
755+ if !consumed {
756+ send_buf.extend_from_slice(&pending_command);
757+ }
758+ pending_command.clear();
759+ command_mode = false;
760+ }
761+ }
687762 if raw_guard.lock().unwrap().is_none() {
688763 continue;
689764 }
690− if input_tx.send(VmInput::Bytes(bytes.to_vec())).is_err() {
765+ if !send_buf.is_empty() && input_tx.send(VmInput::Bytes(send_buf)).is_err()
766+ {
691767 break;
692768 }
693769 }
@@ -753,6 +829,14 @@ pub fn spawn_vm_io(
753829 }
754830}
755831
832+pub fn spawn_vm_io(
833+ output_monitor: Arc<OutputMonitor>,
834+ vm_output_fd: OwnedFd,
835+ vm_input_fd: OwnedFd,
836+) -> IoContext {
837+ spawn_vm_io_with_line_handler(output_monitor, vm_output_fd, vm_input_fd, |_| false)
838+}
839+
756840impl IoContext {
757841 pub fn shutdown(self) {
758842 let _ = self.input_tx.send(VmInput::Shutdown);
@@ -954,13 +1038,17 @@ fn spawn_login_actions_thread(
9541038 })
9551039}
9561040
957−fn run_vm(
1041+fn run_vm_with_io<F>(
9581042 disk_path: &Path,
9591043 login_actions: &[LoginAction],
9601044 directory_shares: &[DirectoryShare],
9611045 cpu_count: usize,
9621046 ram_bytes: u64,
963−) -> Result<(), Box<dyn std::error::Error>> {
1047+ io_handler: F,
1048+) -> Result<(), Box<dyn std::error::Error>>
1049+where
1050+ F: FnOnce(Arc<OutputMonitor>, OwnedFd, OwnedFd) -> IoContext,
1051+{
9641052 let (vm_reads_from, we_write_to) = create_pipe();
9651053 let (we_read_from, vm_writes_to) = create_pipe();
9661054
@@ -1021,7 +1109,7 @@ fn run_vm(
10211109 println!("VM booting...");
10221110
10231111 let output_monitor = Arc::new(OutputMonitor::default());
1024− let io_ctx = spawn_vm_io(output_monitor.clone(), we_read_from, we_write_to);
1112+ let io_ctx = io_handler(output_monitor.clone(), we_read_from, we_write_to);
10251113
10261114 let mut all_login_actions = vec![
10271115 Expect {
@@ -1112,6 +1200,23 @@ fn run_vm(
11121200 exit_result
11131201}
11141202
1203+fn run_vm(
1204+ disk_path: &Path,
1205+ login_actions: &[LoginAction],
1206+ directory_shares: &[DirectoryShare],
1207+ cpu_count: usize,
1208+ ram_bytes: u64,
1209+) -> Result<(), Box<dyn std::error::Error>> {
1210+ run_vm_with_io(
1211+ disk_path,
1212+ login_actions,
1213+ directory_shares,
1214+ cpu_count,
1215+ ram_bytes,
1216+ spawn_vm_io,
1217+ )
1218+}
1219+
11151220fn nsurl_from_path(path: &Path) -> Result<Retained<NSURL>, Box<dyn std::error::Error>> {
11161221 let abs_path = if path.is_absolute() {
11171222 path.to_path_buf()