diff --git a/docs/implementations.md b/docs/implementations.md index 62cd99b..8eb26e5 100644 --- a/docs/implementations.md +++ b/docs/implementations.md @@ -88,4 +88,14 @@ They are also stored per project, in `vibebox.toml` - use `:explain` to display: - mounts: host_path → guest_path, ro/rw - network: mode (allowlist/blocklist) and entries - - storage: paths to vibebox.toml and .vibebox/ (relative from the project_dir) \ No newline at end of file + - storage: paths to vibebox.toml and .vibebox/ (relative from the project_dir) + +## Connection + +### SSH + +- In Project cache, generate and store ssh pair +- In provisioning, install and enable openssh-server in VM +- Mount ssh pair to VM when starting up +- get ipv4 address of VM, store it to project cache +- and connect to VM via ssh with ip and ssh key \ No newline at end of file diff --git a/docs/tasks.md b/docs/tasks.md index 7b69425..b7072f4 100644 --- a/docs/tasks.md +++ b/docs/tasks.md @@ -29,10 +29,13 @@ ## TUI -1. [ ] Fix the terminal component height issue. -2. [ ] Fix the input field that does not expand its height (currently, it just roll the text horizontally). The +1. [x] Fix the terminal component height issue. +2. [x] Fix the input field that does not expand its height (currently, it just roll the text horizontally). The inputfield it should not be scrollable. ## Integration -1. [ ] Integrate VM and SessionManager together. +1. [x] Wire up the vm and tui. +2. [ ] Use ssh to connect to vm. +3. [ ] wire up SessionManager. +4. [ ] VM should be separated by per-session VM daemon process (only accepts if to shutdown vm and itself). diff --git a/src/bin/vibebox-cli.rs b/src/bin/vibebox-cli.rs new file mode 100644 index 0000000..5c27167 --- /dev/null +++ b/src/bin/vibebox-cli.rs @@ -0,0 +1,49 @@ +use std::{ + env, + io::{self, Write}, + sync::{Arc, Mutex}, +}; + +use color_eyre::Result; + +use vibebox::tui::{AppState, VmInfo}; +use vibebox::{tui, vm}; + +fn main() -> Result<()> { + color_eyre::install()?; + + let args = vm::parse_cli().map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + if args.version() { + vm::print_version(); + return Ok(()); + } + if args.help() { + vm::print_help(); + return Ok(()); + } + + let vm_info = VmInfo { + version: env!("CARGO_PKG_VERSION").to_string(), + max_memory_mb: args.ram_mb(), + cpu_cores: args.cpu_count(), + }; + let cwd = env::current_dir().map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + let app = Arc::new(Mutex::new(AppState::new(cwd, vm_info))); + + { + let mut locked = app.lock().expect("app state poisoned"); + tui::render_tui_once(&mut locked)?; + } + { + let mut stdout = io::stdout().lock(); + writeln!(stdout)?; + stdout.flush()?; + } + + vm::run_with_args(args, |output_monitor, vm_output_fd, vm_input_fd| { + tui::passthrough_vm_io(app.clone(), output_monitor, vm_output_fd, vm_input_fd) + }) + .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + + Ok(()) +} diff --git a/src/bin/vibebox-tui.rs b/src/bin/vibebox-tui.rs deleted file mode 100644 index b318c8f..0000000 --- a/src/bin/vibebox-tui.rs +++ /dev/null @@ -1,247 +0,0 @@ -use std::{ - env, - ffi::OsString, - io::{self, Read, Write}, - path::PathBuf, -}; - -use color_eyre::Result; -use lexopt::prelude::*; - -#[path = "../tui.rs"] -mod tui; - -use tui::{AppState, VmInfo}; - -#[derive(Debug, Clone, PartialEq, Eq)] -struct TuiConfig { - cwd: PathBuf, - vm_version: String, - max_memory_mb: u64, - cpu_cores: usize, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -enum TuiCommand { - Run(TuiConfig), - Help, - Version, -} - -#[derive(Debug, thiserror::Error)] -enum CliError { - #[error("{0}")] - Message(String), - #[error(transparent)] - Lexopt(#[from] lexopt::Error), - #[error(transparent)] - Io(#[from] std::io::Error), -} - -fn main() -> Result<()> { - color_eyre::install()?; - - let command = parse_args(env::args_os())?; - match command { - TuiCommand::Help => { - print_help(); - } - TuiCommand::Version => { - println!("vibebox-tui {}", env!("CARGO_PKG_VERSION")); - } - TuiCommand::Run(config) => { - let vm_info = VmInfo { - version: config.vm_version, - max_memory_mb: config.max_memory_mb, - cpu_cores: config.cpu_cores, - }; - let mut app = AppState::new(config.cwd, vm_info); - tui::render_tui_once(&mut app)?; - { - let mut stdout = io::stdout().lock(); - writeln!(stdout)?; - stdout.flush()?; - } - passthrough_stdio(&mut app)?; - } - } - - Ok(()) -} - -fn passthrough_stdio(app: &mut AppState) -> io::Result<()> { - let mut stdin = io::stdin().lock(); - let mut buf = [0u8; 8192]; - let mut line_buf: Vec = Vec::new(); - loop { - let n = stdin.read(&mut buf)?; - if n == 0 { - break; - } - for &b in &buf[..n] { - line_buf.push(b); - if b == b'\n' { - let line = String::from_utf8_lossy(&line_buf); - let trimmed = line.trim_end_matches(&['\r', '\n'][..]); - if trimmed == ":help" { - let _ = tui::render_commands_component(app); - } - line_buf.clear(); - } - } - } - Ok(()) -} - -fn print_help() { - println!( - "vibebox-tui\n\nUsage:\n vibebox-tui [options]\n\nOptions:\n --help, -h Show this help\n --version Show version\n --cwd Working directory for the session header\n --vm-version VM version string for the header\n --max-memory Max memory in MB (default 2048)\n --cpu-cores CPU core count (default 2)\n" - ); -} - -fn parse_args(args: I) -> Result -where - I: IntoIterator, -{ - fn os_to_string(value: OsString, flag: &str) -> Result { - value - .into_string() - .map_err(|_| CliError::Message(format!("{flag} expects valid UTF-8"))) - } - - let mut parser = lexopt::Parser::from_iter(args); - let mut cwd: Option = None; - let mut vm_version = env!("CARGO_PKG_VERSION").to_string(); - let mut max_memory_mb: u64 = 2048; - let mut cpu_cores: usize = 2; - - while let Some(arg) = parser.next()? { - match arg { - Long("help") | Short('h') => return Ok(TuiCommand::Help), - Long("version") => return Ok(TuiCommand::Version), - Long("cwd") => { - let value = os_to_string(parser.value()?, "--cwd")?; - cwd = Some(PathBuf::from(value)); - } - Long("vm-version") => { - vm_version = os_to_string(parser.value()?, "--vm-version")?; - } - Long("max-memory") => { - let value: u64 = os_to_string(parser.value()?, "--max-memory")? - .parse() - .map_err(|_| { - CliError::Message("--max-memory expects an integer".to_string()) - })?; - if value == 0 { - return Err(CliError::Message("--max-memory must be >= 1".to_string())); - } - max_memory_mb = value; - } - Long("cpu-cores") => { - let value: usize = os_to_string(parser.value()?, "--cpu-cores")? - .parse() - .map_err(|_| CliError::Message("--cpu-cores expects an integer".to_string()))?; - if value == 0 { - return Err(CliError::Message("--cpu-cores must be >= 1".to_string())); - } - cpu_cores = value; - } - Value(value) => { - return Err(CliError::Message(format!( - "unexpected argument: {}", - value.to_string_lossy() - ))); - } - _ => return Err(CliError::Message(arg.unexpected().to_string())), - } - } - - let cwd = match cwd { - Some(dir) => dir, - None => env::current_dir()?, - }; - - Ok(TuiCommand::Run(TuiConfig { - cwd, - vm_version, - max_memory_mb, - cpu_cores, - })) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn parse_from(args: &[&str]) -> Result { - let mut argv = vec![OsString::from("vibebox-tui")]; - argv.extend(args.iter().map(OsString::from)); - parse_args(argv) - } - - #[test] - fn parse_help_short_circuit() { - let command = parse_from(&["--help"]).unwrap(); - assert!(matches!(command, TuiCommand::Help)); - } - - #[test] - fn parse_version_short_circuit() { - let command = parse_from(&["--version"]).unwrap(); - assert!(matches!(command, TuiCommand::Version)); - } - - #[test] - fn parse_defaults() { - let command = parse_from(&[]).unwrap(); - let TuiCommand::Run(config) = command else { - panic!("expected run command"); - }; - - assert_eq!(config.vm_version, env!("CARGO_PKG_VERSION")); - assert_eq!(config.max_memory_mb, 2048); - assert_eq!(config.cpu_cores, 2); - } - - #[test] - fn parse_overrides() { - let command = parse_from(&[ - "--cwd", - "/tmp", - "--vm-version", - "13.1", - "--max-memory", - "4096", - "--cpu-cores", - "4", - ]) - .unwrap(); - - let TuiCommand::Run(config) = command else { - panic!("expected run command"); - }; - - assert_eq!(config.cwd, PathBuf::from("/tmp")); - assert_eq!(config.vm_version, "13.1"); - assert_eq!(config.max_memory_mb, 4096); - assert_eq!(config.cpu_cores, 4); - } - - #[test] - fn parse_rejects_zero_cpu() { - let err = parse_from(&["--cpu-cores", "0"]).unwrap_err(); - assert!(err.to_string().contains("cpu-cores")); - } - - #[test] - fn parse_rejects_zero_memory() { - let err = parse_from(&["--max-memory", "0"]).unwrap_err(); - assert!(err.to_string().contains("max-memory")); - } - - #[test] - fn parse_rejects_unknown_argument() { - let err = parse_from(&["--unknown"]).unwrap_err(); - assert!(!err.to_string().is_empty()); - } -} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..4155675 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,5 @@ +pub mod session_manager; +pub mod tui; +pub mod vm; + +pub use session_manager::{SessionError, SessionManager, SessionRecord}; diff --git a/src/tui.rs b/src/tui.rs index 8e4a70a..af026f9 100644 --- a/src/tui.rs +++ b/src/tui.rs @@ -1,6 +1,8 @@ use std::{ io::{self, Write}, + os::unix::io::OwnedFd, path::PathBuf, + sync::{Arc, Mutex}, }; use color_eyre::Result; @@ -21,6 +23,8 @@ use ratatui::{ widgets::{Block, Borders, List, ListItem, Paragraph, Widget}, }; +use crate::vm; + const ASCII_BANNER: [&str; 7] = [ "██╗ ██╗██╗██████╗ ███████╗██████╗ ██████╗ ██╗ ██╗", "██║ ██║██║██╔══██╗██╔════╝██╔══██╗██╔═══██╗╚██╗██╔╝", @@ -170,6 +174,23 @@ pub fn render_commands_component(app: &mut AppState) -> Result<()> { Ok(()) } +pub fn passthrough_vm_io( + app: Arc>, + output_monitor: Arc, + vm_output_fd: OwnedFd, + vm_input_fd: OwnedFd, +) -> vm::IoContext { + vm::spawn_vm_io_with_line_handler(output_monitor, vm_output_fd, vm_input_fd, move |line| { + if line == ":help" { + if let Ok(mut locked) = app.lock() { + let _ = render_commands_component(&mut locked); + } + return true; + } + false + }) +} + fn render_static_buffer(app: &mut AppState, width: u16) -> Buffer { let layout = compute_page_layout(app, width); let content_height = layout.total_height.max(1); diff --git a/src/main.rs b/src/vm.rs similarity index 91% rename from src/main.rs rename to src/vm.rs index 687c732..0e021fa 100644 --- a/src/main.rs +++ b/src/vm.rs @@ -1,7 +1,3 @@ -mod session_manager; - -pub use session_manager::{SessionError, SessionManager, SessionRecord}; - use std::{ env, ffi::OsString, @@ -119,41 +115,26 @@ impl DirectoryShare { } } -fn main() -> Result<(), Box> { +pub fn run_cli() -> Result<(), Box> { let args = parse_cli()?; - if args.version { - println!("Vibe"); - println!("https://github.com/lynaghk/vibe/"); - println!("Git SHA: {}", env!("GIT_SHA")); - std::process::exit(0); + if args.version() { + print_version(); + return Ok(()); } - if args.help { - println!( - "Vibe is a quick way to spin up a Linux virtual machine on Mac to sandbox LLM agents. - -vibe [OPTIONS] [disk-image.raw] - -Options - - --help Print this help message. - --version Print the version (commit SHA). - --no-default-mounts Disable all default mounts. - --mount host-path:guest-path[:read-only | :read-write] Mount `host-path` inside VM at `guest-path`. - Defaults to read-write. - Errors if host-path does not exist. - --cpus Number of virtual CPUs (default {DEFAULT_CPU_COUNT}). - --ram RAM size in megabytes (default {DEFAULT_RAM_MB}). - --script Run script in VM. - --send Type `some-command` followed by newline into the VM. - --expect [timeout-seconds] Wait for `string` to appear in console output before executing next `--script` or `--send`. - If `string` does not appear within timeout (default 30 seconds), shutdown VM with error. -" - ); - std::process::exit(0); + if args.help() { + print_help(); + return Ok(()); } + run_with_args(args, spawn_vm_io) +} + +pub fn run_with_args(args: CliArgs, io_handler: F) -> Result<(), Box> +where + F: FnOnce(Arc, OwnedFd, OwnedFd) -> IoContext, +{ ensure_signed(); let project_root = env::current_dir()?; @@ -239,6 +220,11 @@ Options ), DirectoryShare::new(home.join(".codex"), "/root/.codex".into(), false), DirectoryShare::new(home.join(".claude"), "/root/.claude".into(), false), + DirectoryShare::new( + "/Users/zhangjie/Documents/Code/CompletePrograms/vibebox/.ssh".into(), + "/root/.ssh".into(), + true, + ), ] .into_iter() .flatten() @@ -258,16 +244,47 @@ Options // Any user-provided login actions must come after our system ones login_actions.extend(args.login_actions); - run_vm( + run_vm_with_io( &disk_path, &login_actions, &directory_shares[..], args.cpu_count, args.ram_bytes, + io_handler, ) } -struct CliArgs { +pub fn print_help() { + println!( + "Vibe is a quick way to spin up a Linux virtual machine on Mac to sandbox LLM agents. + +vibe [OPTIONS] [disk-image.raw] + +Options + + --help Print this help message. + --version Print the version (commit SHA). + --no-default-mounts Disable all default mounts. + --mount host-path:guest-path[:read-only | :read-write] Mount `host-path` inside VM at `guest-path`. + Defaults to read-write. + Errors if host-path does not exist. + --cpus Number of virtual CPUs (default {DEFAULT_CPU_COUNT}). + --ram RAM size in megabytes (default {DEFAULT_RAM_MB}). + --script Run script in VM. + --send Type `some-command` followed by newline into the VM. + --expect [timeout-seconds] Wait for `string` to appear in console output before executing next `--script` or `--send`. + If `string` does not appear within timeout (default 30 seconds), shutdown VM with error. +" + ); +} + +pub fn print_version() { + println!("Vibe"); + println!("https://github.com/lynaghk/vibe/"); + println!("Git SHA: {}", env!("GIT_SHA")); +} + +pub struct CliArgs { disk: Option, version: bool, help: bool, @@ -278,14 +295,43 @@ struct CliArgs { ram_bytes: u64, } -fn parse_cli() -> Result> { +impl CliArgs { + pub fn version(&self) -> bool { + self.version + } + + pub fn help(&self) -> bool { + self.help + } + + pub fn cpu_count(&self) -> usize { + self.cpu_count + } + + pub fn ram_bytes(&self) -> u64 { + self.ram_bytes + } + + pub fn ram_mb(&self) -> u64 { + self.ram_bytes / BYTES_PER_MB + } +} + +pub fn parse_cli() -> Result> { + parse_cli_from(env::args_os()) +} + +pub fn parse_cli_from(args: I) -> Result> +where + I: IntoIterator, +{ fn os_to_string(value: OsString, flag: &str) -> Result> { value .into_string() .map_err(|_| format!("{flag} expects valid UTF-8").into()) } - let mut parser = lexopt::Parser::from_env(); + let mut parser = lexopt::Parser::from_iter(args); let mut disk = None; let mut version = false; let mut help = false; @@ -621,11 +667,15 @@ pub fn create_pipe() -> (OwnedFd, OwnedFd) { (read_stream.into(), write_stream.into()) } -pub fn spawn_vm_io( +pub fn spawn_vm_io_with_line_handler( output_monitor: Arc, vm_output_fd: OwnedFd, vm_input_fd: OwnedFd, -) -> IoContext { + mut on_line: F, +) -> IoContext +where + F: FnMut(&str) -> bool + ::std::marker::Send + 'static, +{ let (input_tx, input_rx): (Sender, Receiver) = mpsc::channel(); // 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( move || { let mut buf = [0u8; 1024]; + let mut pending_command: Vec = Vec::new(); + let mut command_mode = false; loop { match poll_with_wakeup(libc::STDIN_FILENO, wakeup_read.as_raw_fd(), &mut buf) { PollResult::Shutdown | PollResult::Error => break, PollResult::Spurious => continue, PollResult::Ready(bytes) => { + let mut send_buf: Vec = Vec::new(); + for &b in bytes { + if pending_command.is_empty() && !command_mode && b == b':' { + command_mode = true; + } + + if command_mode { + pending_command.push(b); + } else { + send_buf.push(b); + } + + if b == b'\n' && command_mode { + let line = String::from_utf8_lossy(&pending_command); + let trimmed = line.trim_end_matches(&['\r', '\n'][..]); + let consumed = on_line(trimmed); + if !consumed { + send_buf.extend_from_slice(&pending_command); + } + pending_command.clear(); + command_mode = false; + } + } if raw_guard.lock().unwrap().is_none() { continue; } - if input_tx.send(VmInput::Bytes(bytes.to_vec())).is_err() { + if !send_buf.is_empty() && input_tx.send(VmInput::Bytes(send_buf)).is_err() + { break; } } @@ -753,6 +829,14 @@ pub fn spawn_vm_io( } } +pub fn spawn_vm_io( + output_monitor: Arc, + vm_output_fd: OwnedFd, + vm_input_fd: OwnedFd, +) -> IoContext { + spawn_vm_io_with_line_handler(output_monitor, vm_output_fd, vm_input_fd, |_| false) +} + impl IoContext { pub fn shutdown(self) { let _ = self.input_tx.send(VmInput::Shutdown); @@ -954,13 +1038,17 @@ fn spawn_login_actions_thread( }) } -fn run_vm( +fn run_vm_with_io( disk_path: &Path, login_actions: &[LoginAction], directory_shares: &[DirectoryShare], cpu_count: usize, ram_bytes: u64, -) -> Result<(), Box> { + io_handler: F, +) -> Result<(), Box> +where + F: FnOnce(Arc, OwnedFd, OwnedFd) -> IoContext, +{ let (vm_reads_from, we_write_to) = create_pipe(); let (we_read_from, vm_writes_to) = create_pipe(); @@ -1021,7 +1109,7 @@ fn run_vm( println!("VM booting..."); let output_monitor = Arc::new(OutputMonitor::default()); - let io_ctx = spawn_vm_io(output_monitor.clone(), we_read_from, we_write_to); + let io_ctx = io_handler(output_monitor.clone(), we_read_from, we_write_to); let mut all_login_actions = vec![ Expect { @@ -1112,6 +1200,23 @@ fn run_vm( exit_result } +fn run_vm( + disk_path: &Path, + login_actions: &[LoginAction], + directory_shares: &[DirectoryShare], + cpu_count: usize, + ram_bytes: u64, +) -> Result<(), Box> { + run_vm_with_io( + disk_path, + login_actions, + directory_shares, + cpu_count, + ram_bytes, + spawn_vm_io, + ) +} + fn nsurl_from_path(path: &Path) -> Result, Box> { let abs_path = if path.is_absolute() { path.to_path_buf()