From ca33633d28b60ebf0843293edfd61d9108063b00 Mon Sep 17 00:00:00 2001 From: robcholz <84130577+robcholz@users.noreply.github.com> Date: Sat, 7 Feb 2026 00:22:24 -0500 Subject: [PATCH] feat: allow multi vibebox to connect to the same vm. --- Cargo.lock | 76 +++++ Cargo.toml | 2 + docs/arch.mermaid | 14 + docs/tasks.md | 19 +- src/bin/vibebox-cli.rs | 67 ++++- src/bin/vibebox-supervisor.rs | 43 +++ src/instance.rs | 180 ++++++++---- src/lib.rs | 1 + src/vm.rs | 3 + src/vm_manager.rs | 538 ++++++++++++++++++++++++++++++++++ 10 files changed, 884 insertions(+), 59 deletions(-) create mode 100644 docs/arch.mermaid create mode 100644 src/bin/vibebox-supervisor.rs create mode 100644 src/vm_manager.rs diff --git a/Cargo.lock b/Cargo.lock index 2e7d84f..eb00dea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,15 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -506,6 +515,15 @@ dependencies = [ "hashbrown 0.15.5", ] +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "memchr" version = "2.7.6" @@ -533,6 +551,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.59.0", +] + [[package]] name = "num-conv" version = "0.2.0" @@ -726,6 +753,23 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a96887878f22d7bad8a3b6dc5b7440e0ada9a245242924394987b21cf2210a4c" + [[package]] name = "rustc-demangle" version = "0.1.27" @@ -1069,9 +1113,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.36" @@ -1092,15 +1148,33 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + [[package]] name = "tracing-subscriber" version = "0.3.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f30143827ddab0d256fd843b7a66d164e9f271cfa0dde49142c5ca0ca291f1e" dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", "sharded-slab", + "smallvec", "thread_local", + "tracing", "tracing-core", + "tracing-log", ] [[package]] @@ -1187,6 +1261,8 @@ dependencies = [ "time", "tokio", "toml", + "tracing", + "tracing-subscriber", "tui-textarea", "unicode-width 0.2.0", "uuid", diff --git a/Cargo.toml b/Cargo.toml index 3c0fb6e..f0e8f3a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,3 +34,5 @@ ratatui = { version = "0.29.0", features = ["unstable-rendered-line-info"] } tokio = { version = "1.40.0", features = ["full"] } tui-textarea = { version = "0.7.0", default-features = false, features = ["ratatui"] } unicode-width = "0.2" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/docs/arch.mermaid b/docs/arch.mermaid new file mode 100644 index 0000000..0aa6544 --- /dev/null +++ b/docs/arch.mermaid @@ -0,0 +1,14 @@ +flowchart TD + VB[Vibebox] --> SM[Session Manager] + + SM --> VMM1[VM Manager #1] + SM --> VMM2[VM Manager #2] + SM --> VMM3[VM Manager #N] + + VMM1 --> VMI1[VM Instance #1] + VMM2 --> VMI2[VM Instance #2] + VMM3 --> VMIN[VM Instance #N] + + VMI1 --> VM1[VM #1] + VMI2 --> VM2[VM #2] + VMIN --> VMN[VM #N] \ No newline at end of file diff --git a/docs/tasks.md b/docs/tasks.md index 78b9bfd..2d57cec 100644 --- a/docs/tasks.md +++ b/docs/tasks.md @@ -33,9 +33,22 @@ 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 +## Stage 1 1. [x] Wire up the vm and tui. 2. [x] Use ssh to connect to vm. -3. [ ] wire up SessionManager. -4. [ ] VM should be separated by a per-session VM daemon process (only accepts if to shut down vm and itself). +3. [x] allow multi vibebox to connect to the same vm. +4. [ ] use vm.lock to ensure process concurrency safety. +5. [ ] wire up SessionManager. +6. [ ] VM should be separated by a per-session VM daemon process (only accepts if to shut down vm and itself). +7. [ ] setup vibebox commands +8. [ ] setup cli commands. +9. [ ] fix ui overlap. + +## Publish + +1. [ ] write the docs +2. [ ] setup quick link. +3. [ ] setup website. + +[ ] diff --git a/src/bin/vibebox-cli.rs b/src/bin/vibebox-cli.rs index 6b14c60..e68dd38 100644 --- a/src/bin/vibebox-cli.rs +++ b/src/bin/vibebox-cli.rs @@ -1,18 +1,49 @@ use std::{ env, - io::{self, Write}, + ffi::OsString, + fs, + io::{self, IsTerminal, Write}, + path::Path, sync::{Arc, Mutex}, }; use color_eyre::Result; +use serde::Deserialize; +use tracing_subscriber::EnvFilter; use vibebox::tui::{AppState, VmInfo}; -use vibebox::{instance, tui, vm}; +use vibebox::{instance, tui, vm, vm_manager}; + +const DEFAULT_AUTO_SHUTDOWN_MS: u64 = 3000; + +#[derive(Debug, Default, Deserialize)] +struct ProjectConfig { + auto_shutdown_ms: Option, +} fn main() -> Result<()> { + init_tracing(); color_eyre::install()?; + let raw_args: Vec = env::args_os().collect(); + + if env::var("VIBEBOX_VM_MANAGER").as_deref() == Ok("1") { + tracing::info!("starting vm manager mode"); + let args = vm::parse_cli().map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + let auto_shutdown_ms = env::var("VIBEBOX_AUTO_SHUTDOWN_MS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_AUTO_SHUTDOWN_MS); + tracing::info!(auto_shutdown_ms, "vm manager config"); + if let Err(err) = vm_manager::run_manager(args, auto_shutdown_ms) { + tracing::error!(error = %err, "vm manager exited"); + return Err(color_eyre::eyre::eyre!(err.to_string())); + } + return Ok(()); + } + let args = vm::parse_cli().map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + tracing::debug!("parsed cli args"); if args.version() { vm::print_version(); return Ok(()); @@ -22,13 +53,16 @@ fn main() -> Result<()> { return Ok(()); } + vm::ensure_signed(); + 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))); + tracing::info!(cwd = %cwd.display(), "starting vibebox cli"); + let app = Arc::new(Mutex::new(AppState::new(cwd.clone(), vm_info))); { let mut locked = app.lock().expect("app state poisoned"); @@ -40,8 +74,33 @@ fn main() -> Result<()> { stdout.flush()?; } - instance::run_with_ssh(args, app.clone()) + let auto_shutdown_ms = load_auto_shutdown_ms(&cwd)?; + tracing::info!(auto_shutdown_ms, "auto shutdown config"); + let manager_conn = vm_manager::ensure_manager(&raw_args, auto_shutdown_ms) .map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + instance::run_with_ssh(manager_conn).map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + Ok(()) } + +fn load_auto_shutdown_ms(project_root: &Path) -> Result { + let path = project_root.join("vibebox.toml"); + let config = match fs::read_to_string(&path) { + Ok(raw) => toml::from_str::(&raw)?, + Err(err) if err.kind() == io::ErrorKind::NotFound => ProjectConfig::default(), + Err(err) => return Err(err.into()), + }; + Ok(config.auto_shutdown_ms.unwrap_or(DEFAULT_AUTO_SHUTDOWN_MS)) +} + +fn init_tracing() { + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); + let ansi = std::io::stderr().is_terminal() && env::var("VIBEBOX_LOG_NO_COLOR").is_err(); + let _ = tracing_subscriber::fmt() + .with_env_filter(filter) + .with_target(false) + .with_ansi(ansi) + .with_writer(std::io::stderr) + .try_init(); +} diff --git a/src/bin/vibebox-supervisor.rs b/src/bin/vibebox-supervisor.rs new file mode 100644 index 0000000..acc6401 --- /dev/null +++ b/src/bin/vibebox-supervisor.rs @@ -0,0 +1,43 @@ +use std::{ + env, + io::{self, IsTerminal}, +}; + +use color_eyre::Result; +use tracing_subscriber::EnvFilter; + +use vibebox::{vm, vm_manager}; + +const DEFAULT_AUTO_SHUTDOWN_MS: u64 = 3000; + +fn main() -> Result<()> { + init_tracing(); + color_eyre::install()?; + + tracing::info!("starting vm supervisor"); + let args = vm::parse_cli().map_err(|err| color_eyre::eyre::eyre!(err.to_string()))?; + let auto_shutdown_ms = env::var("VIBEBOX_AUTO_SHUTDOWN_MS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_AUTO_SHUTDOWN_MS); + tracing::info!(auto_shutdown_ms, "vm supervisor config"); + + if let Err(err) = vm_manager::run_manager(args, auto_shutdown_ms) { + tracing::error!(error = %err, "vm supervisor exited"); + return Err(color_eyre::eyre::eyre!(err.to_string())); + } + tracing::info!("vm supervisor exited"); + + Ok(()) +} + +fn init_tracing() { + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); + let ansi = std::io::stderr().is_terminal() && env::var("VIBEBOX_LOG_NO_COLOR").is_err(); + let _ = tracing_subscriber::fmt() + .with_env_filter(filter) + .with_target(false) + .with_ansi(ansi) + .with_writer(io::stderr) + .try_init(); +} diff --git a/src/instance.rs b/src/instance.rs index 0890f94..7f3e062 100644 --- a/src/instance.rs +++ b/src/instance.rs @@ -2,7 +2,7 @@ use std::{ env, fs, io::{self, Write}, net::{SocketAddr, TcpStream}, - os::unix::{fs::PermissionsExt, io::OwnedFd}, + os::unix::{fs::PermissionsExt, io::OwnedFd, net::UnixStream}, path::{Path, PathBuf}, process::{Command, Stdio}, sync::{ @@ -10,6 +10,7 @@ use std::{ atomic::{AtomicBool, Ordering}, }, thread, + time::{Duration, Instant}, }; use serde::{Deserialize, Serialize}; @@ -17,13 +18,14 @@ use std::sync::mpsc::Sender; use uuid::Uuid; use crate::{ - session_manager::{GLOBAL_DIR_NAME, INSTANCE_DIR_NAME}, + session_manager::INSTANCE_DIR_NAME, tui::{self, AppState}, - vm::{self, DirectoryShare, LoginAction, VmInput}, + vm::{self, LoginAction, VmInput}, }; const INSTANCE_TOML: &str = "instance.toml"; const SSH_KEY_NAME: &str = "ssh_key"; +#[allow(dead_code)] const SERIAL_LOG_NAME: &str = "serial.log"; const DEFAULT_SSH_USER: &str = "vibecoder"; const SSH_CONNECT_RETRIES: usize = 30; @@ -31,76 +33,49 @@ const SSH_CONNECT_DELAY_MS: u64 = 500; const SSH_SETUP_SCRIPT: &str = include_str!("ssh.sh"); #[derive(Debug, Clone, Serialize, Deserialize)] -struct InstanceConfig { +pub(crate) struct InstanceConfig { #[serde(default = "default_ssh_user")] ssh_user: String, #[serde(default)] sudo_password: String, #[serde(default)] - vm_ipv4: Option, + pub(crate) vm_ipv4: Option, } fn default_ssh_user() -> String { DEFAULT_SSH_USER.to_string() } -pub fn run_with_ssh( - args: vm::CliArgs, - app: Arc>, -) -> Result<(), Box> { +pub fn run_with_ssh(manager_conn: UnixStream) -> Result<(), Box> { let project_root = env::current_dir()?; - let project_name = project_root - .file_name() - .ok_or("Project directory has no name")? - .to_string_lossy() - .into_owned(); + tracing::info!(root = %project_root.display(), "starting ssh session"); let instance_dir = ensure_instance_dir(&project_root)?; + tracing::debug!(instance_dir = %instance_dir.display(), "instance dir ready"); let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?; - let mut config = load_or_create_instance_config(&instance_dir)?; - // Clear cached IP to avoid reusing a stale address on startup. - if config.vm_ipv4.is_some() { - config.vm_ipv4 = None; - write_instance_config(&instance_dir.join(INSTANCE_TOML), &config)?; - } - let config = Arc::new(Mutex::new(config)); + let config = load_or_create_instance_config(&instance_dir)?; + let ssh_user = config.ssh_user.clone(); + tracing::debug!(ssh_user = %ssh_user, "loaded instance config"); - let ssh_guest_dir = format!("/root/{}", GLOBAL_DIR_NAME); + let _manager_conn = manager_conn; + tracing::debug!("waiting for vm ipv4"); + wait_for_vm_ipv4(&instance_dir, Duration::from_secs(120))?; - let extra_shares = vec![DirectoryShare::new( - instance_dir.clone(), - ssh_guest_dir.clone().into(), - true, - )?]; + let ip = load_or_create_instance_config(&instance_dir)? + .vm_ipv4 + .ok_or("VM IPv4 not available")?; + tracing::info!(ip = %ip, "vm ipv4 ready"); - let extra_login_actions = - build_ssh_login_actions(&config, &project_name, ssh_guest_dir.as_str(), SSH_KEY_NAME); - - vm::run_with_args_and_extras( - args, - |output_monitor, vm_output_fd, vm_input_fd| { - spawn_ssh_io( - app.clone(), - config.clone(), - instance_dir.clone(), - ssh_key.clone(), - output_monitor, - vm_output_fd, - vm_input_fd, - ) - }, - extra_login_actions, - extra_shares, - ) + run_ssh_session(ssh_key, ssh_user, ip) } -fn ensure_instance_dir(project_root: &Path) -> Result { +pub(crate) fn ensure_instance_dir(project_root: &Path) -> Result { let instance_dir = project_root.join(INSTANCE_DIR_NAME); fs::create_dir_all(&instance_dir)?; Ok(instance_dir) } -fn ensure_ssh_keypair( +pub(crate) fn ensure_ssh_keypair( instance_dir: &Path, ) -> Result<(PathBuf, PathBuf), Box> { let private_key = instance_dir.join(SSH_KEY_NAME); @@ -143,7 +118,7 @@ fn ensure_ssh_keypair( Ok((private_key, public_key)) } -fn load_or_create_instance_config( +pub(crate) fn load_or_create_instance_config( instance_dir: &Path, ) -> Result> { let config_path = instance_dir.join(INSTANCE_TOML); @@ -176,7 +151,7 @@ fn load_or_create_instance_config( Ok(config) } -fn write_instance_config( +pub(crate) fn write_instance_config( path: &Path, config: &InstanceConfig, ) -> Result<(), Box> { @@ -190,7 +165,7 @@ fn generate_password() -> String { Uuid::now_v7().simple().to_string() } -fn extract_ipv4(line: &str) -> Option { +pub(crate) fn extract_ipv4(line: &str) -> Option { let mut current = String::new(); let mut best: Option = None; @@ -209,6 +184,106 @@ fn extract_ipv4(line: &str) -> Option { best } +fn wait_for_vm_ipv4( + instance_dir: &Path, + timeout: Duration, +) -> Result<(), Box> { + let start = Instant::now(); + loop { + let config = load_or_create_instance_config(instance_dir)?; + if config.vm_ipv4.is_some() { + return Ok(()); + } + if start.elapsed() > timeout { + return Err("Timed out waiting for VM IPv4".into()); + } + if start.elapsed().as_secs() % 5 == 0 { + tracing::debug!("still waiting for vm ipv4"); + } + thread::sleep(Duration::from_millis(200)); + } +} + +fn run_ssh_session( + ssh_key: PathBuf, + ssh_user: String, + ip: String, +) -> Result<(), Box> { + let mut attempts = 0usize; + loop { + attempts += 1; + if !ssh_port_open(&ip) { + tracing::debug!(attempts, "ssh port not open yet"); + eprintln!( + "[vibebox] waiting for ssh port on {ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})" + ); + if attempts >= SSH_CONNECT_RETRIES { + return Err( + format!("ssh port not ready after {SSH_CONNECT_RETRIES} attempts").into(), + ); + } + thread::sleep(Duration::from_millis(SSH_CONNECT_DELAY_MS)); + continue; + } + + eprintln!( + "[vibebox] starting ssh to {ssh_user}@{ip} (attempt {attempts}/{SSH_CONNECT_RETRIES})" + ); + tracing::info!(attempts, user = %ssh_user, ip = %ip, "starting ssh"); + let status = Command::new("ssh") + .args([ + "-i", + ssh_key.to_str().unwrap_or(".vibebox/ssh_key"), + "-o", + "IdentitiesOnly=yes", + "-o", + "StrictHostKeyChecking=no", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "GlobalKnownHostsFile=/dev/null", + "-o", + "PasswordAuthentication=no", + "-o", + "BatchMode=yes", + "-o", + "LogLevel=ERROR", + "-o", + "ConnectTimeout=5", + ]) + .env_remove("LC_CTYPE") + .env_remove("LC_ALL") + .env_remove("LANG") + .arg(format!("{ssh_user}@{ip}")) + .stdin(Stdio::inherit()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status(); + + match status { + Ok(status) if status.success() => { + eprintln!("[vibebox] ssh exited with status: {status}"); + break; + } + Ok(status) if status.code() == Some(255) => { + eprintln!("[vibebox] ssh connection failed: {status}"); + if attempts >= SSH_CONNECT_RETRIES { + return Err(format!("ssh failed after {SSH_CONNECT_RETRIES} attempts").into()); + } + thread::sleep(Duration::from_millis(500)); + } + Ok(status) => { + return Err(format!("ssh exited with status: {status}").into()); + } + Err(err) => { + return Err(format!("failed to start ssh: {err}").into()); + } + } + } + + Ok(()) +} + fn is_ipv4_candidate(candidate: &str) -> bool { let parts: Vec<&str> = candidate.split('.').collect(); if parts.len() != 4 { @@ -233,7 +308,7 @@ fn ssh_port_open(ip: &str) -> bool { TcpStream::connect_timeout(&addr, std::time::Duration::from_millis(500)).is_ok() } -fn build_ssh_login_actions( +pub(crate) fn build_ssh_login_actions( config: &Arc>, project_name: &str, guest_dir: &str, @@ -257,6 +332,7 @@ fn build_ssh_login_actions( vec![LoginAction::Send(setup)] } +#[allow(dead_code)] fn spawn_ssh_io( app: Arc>, config: Arc>, diff --git a/src/lib.rs b/src/lib.rs index d089ae0..e1eaad0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,5 +2,6 @@ pub mod instance; pub mod session_manager; pub mod tui; pub mod vm; +pub mod vm_manager; pub use session_manager::{SessionError, SessionManager, SessionRecord}; diff --git a/src/vm.rs b/src/vm.rs index b354d42..336e3dd 100644 --- a/src/vm.rs +++ b/src/vm.rs @@ -1358,6 +1358,9 @@ impl Drop for RawModeGuard { // Ensure the running binary has com.apple.security.virtualization entitlements by checking and, if not, signing and relaunching. pub fn ensure_signed() { + if std::env::var("VIBEBOX_SKIP_CODESIGN").as_deref() == Ok("1") { + return; + } let exe = std::env::current_exe().expect("failed to get current exe path"); let exe_str = exe.to_str().expect("exe path not valid utf-8"); diff --git a/src/vm_manager.rs b/src/vm_manager.rs new file mode 100644 index 0000000..ec93792 --- /dev/null +++ b/src/vm_manager.rs @@ -0,0 +1,538 @@ +use std::{ + env, fs, + io::{Read, Write}, + os::unix::{ + fs::PermissionsExt, + io::AsRawFd, + net::{UnixListener, UnixStream}, + process::CommandExt, + }, + path::{Path, PathBuf}, + process::{Command, Stdio}, + sync::{Arc, Mutex, mpsc}, + thread, + time::{Duration, Instant}, +}; + +use crate::{ + instance::{ + InstanceConfig, build_ssh_login_actions, ensure_instance_dir, ensure_ssh_keypair, + extract_ipv4, load_or_create_instance_config, write_instance_config, + }, + session_manager::GLOBAL_DIR_NAME, + vm::{self, DirectoryShare, LoginAction, VmInput}, +}; + +const VM_MANAGER_SOCKET_NAME: &str = "vm.sock"; + +pub fn ensure_manager( + raw_args: &[std::ffi::OsString], + auto_shutdown_ms: u64, +) -> Result> { + let project_root = env::current_dir()?; + tracing::debug!(root = %project_root.display(), "ensure vm manager"); + let instance_dir = ensure_instance_dir(&project_root)?; + let socket_path = instance_dir.join(VM_MANAGER_SOCKET_NAME); + + if let Ok(stream) = UnixStream::connect(&socket_path) { + send_client_pid(&stream); + tracing::info!(path = %socket_path.display(), "connected to existing vm manager"); + return Ok(stream); + } + + tracing::info!(path = %socket_path.display(), "spawning vm manager"); + spawn_manager_process(raw_args, auto_shutdown_ms, &instance_dir)?; + + let start = Instant::now(); + let timeout = Duration::from_secs(10); + loop { + match UnixStream::connect(&socket_path) { + Ok(stream) => { + send_client_pid(&stream); + tracing::info!(path = %socket_path.display(), "connected to vm manager"); + return Ok(stream); + } + Err(err) => { + tracing::debug!(error = %err, "waiting for vm manager socket"); + if start.elapsed() > timeout { + return Err(format!( + "Timed out waiting for vm manager socket: {} ({})", + socket_path.display(), + err + ) + .into()); + } + thread::sleep(Duration::from_millis(100)); + } + } + } +} + +pub fn run_manager( + args: vm::CliArgs, + auto_shutdown_ms: u64, +) -> Result<(), Box> { + let project_root = env::current_dir()?; + tracing::info!(root = %project_root.display(), "vm manager starting"); + run_manager_with( + &project_root, + args, + auto_shutdown_ms, + &RealVmExecutor, + ManagerOptions { + ensure_signed: true, + detach: true, + prepare_vm: true, + }, + ) +} + +fn spawn_manager_process( + raw_args: &[std::ffi::OsString], + auto_shutdown_ms: u64, + instance_dir: &Path, +) -> Result<(), Box> { + let exe = env::current_exe()?; + let mut supervisor_exe = exe.clone(); + supervisor_exe.set_file_name("vibebox-supervisor"); + let use_supervisor = supervisor_exe.exists(); + let mut cmd = if use_supervisor { + Command::new(supervisor_exe) + } else { + let mut cmd = Command::new(exe); + cmd.arg0("vibebox-supervisor"); + cmd + }; + if raw_args.len() > 1 { + cmd.args(&raw_args[1..]); + } + if !use_supervisor { + cmd.env("VIBEBOX_VM_MANAGER", "1"); + } + cmd.env("VIBEBOX_LOG_NO_COLOR", "1"); + cmd.env("VIBEBOX_AUTO_SHUTDOWN_MS", auto_shutdown_ms.to_string()); + tracing::info!(auto_shutdown_ms, "vm manager process spawn requested"); + let log_path = instance_dir.join("vm_manager.log"); + let log_file = fs::OpenOptions::new() + .create(true) + .append(true) + .open(&log_path) + .ok(); + if let Some(file) = log_file { + let stderr = Stdio::from(file); + cmd.stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(stderr); + } else { + cmd.stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + } + let _child = cmd.spawn()?; + Ok(()) +} + +fn detach_from_terminal() { + unsafe { + libc::setsid(); + } + + if let Ok(devnull) = fs::OpenOptions::new() + .read(true) + .write(true) + .open("/dev/null") + { + let _ = unsafe { libc::dup2(devnull.as_raw_fd(), libc::STDIN_FILENO) }; + let _ = unsafe { libc::dup2(devnull.as_raw_fd(), libc::STDOUT_FILENO) }; + } +} + +fn wait_for_disconnect(mut stream: UnixStream) { + let mut buf = [0u8; 64]; + loop { + match stream.read(&mut buf) { + Ok(0) => break, + Ok(_) => continue, + Err(_) => break, + } + } +} + +fn send_client_pid(stream: &UnixStream) { + let pid = std::process::id(); + let payload = format!("pid={pid}\n"); + if let Ok(mut stream) = stream.try_clone() { + let _ = stream.write_all(payload.as_bytes()); + let _ = stream.flush(); + } +} + +fn read_client_pid(stream: &UnixStream) -> Option { + let mut stream = stream.try_clone().ok()?; + let _ = stream.set_read_timeout(Some(Duration::from_millis(200))); + let mut buf = [0u8; 64]; + let mut len = 0usize; + loop { + match stream.read(&mut buf[len..]) { + Ok(0) => break, + Ok(n) => { + len += n; + if buf[..len].iter().any(|b| *b == b'\n') || len == buf.len() { + break; + } + } + Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => break, + Err(err) if err.kind() == std::io::ErrorKind::TimedOut => break, + Err(_) => break, + } + } + let _ = stream.set_read_timeout(None); + if len == 0 { + return None; + } + let line = String::from_utf8_lossy(&buf[..len]); + let trimmed = line.trim(); + if let Some(value) = trimmed.strip_prefix("pid=") { + value.parse::().ok() + } else { + None + } +} + +fn spawn_manager_io( + config: Arc>, + instance_dir: PathBuf, + output_monitor: Arc, + vm_output_fd: std::os::unix::io::OwnedFd, + vm_input_fd: std::os::unix::io::OwnedFd, +) -> vm::IoContext { + let log_path = instance_dir.join("serial.log"); + let log_file = fs::OpenOptions::new() + .create(true) + .append(true) + .open(&log_path) + .ok() + .map(|file| Arc::new(Mutex::new(file))); + + let instance_path = instance_dir.join("instance.toml"); + let config_for_output = config.clone(); + let log_for_output = log_file.clone(); + let mut line_buf = String::new(); + + let on_output = move |bytes: &[u8]| { + if let Some(log) = &log_for_output { + if let Ok(mut file) = log.lock() { + let _ = file.write_all(bytes); + } + } + + let text = String::from_utf8_lossy(bytes); + line_buf.push_str(&text); + + while let Some(pos) = line_buf.find('\n') { + let mut line = line_buf[..pos].to_string(); + if line.ends_with('\r') { + line.pop(); + } + line_buf.drain(..=pos); + + let cleaned = line.trim_start_matches(|c: char| c == '\r' || c == ' '); + if let Some(pos) = cleaned.find("VIBEBOX_IPV4=") { + let ip_raw = &cleaned[(pos + "VIBEBOX_IPV4=".len())..]; + let ip = extract_ipv4(ip_raw).unwrap_or_default(); + if !ip.is_empty() { + if let Ok(mut cfg) = config_for_output.lock() { + if cfg.vm_ipv4.as_deref() != Some(ip.as_str()) { + cfg.vm_ipv4 = Some(ip.clone()); + let _ = write_instance_config(&instance_path, &cfg); + } + } + } + } + } + }; + + vm::spawn_vm_io_with_hooks( + output_monitor, + vm_output_fd, + vm_input_fd, + vm::IoControl::new(), + |_| false, + on_output, + ) +} + +enum ManagerEvent { + Inc(Option), + Dec(Option), + VmExited(Option), +} + +#[cfg(test)] +mod tests { + use super::*; + use std::{sync::mpsc, time::Duration}; + + #[test] + fn manager_powers_off_after_grace_when_no_refs() { + let _temp = tempfile::Builder::new() + .prefix("vb") + .tempdir_in("/tmp") + .expect("tempdir"); + + let (event_tx, event_rx) = mpsc::channel::(); + let (vm_tx, vm_rx) = mpsc::channel::(); + let vm_input_tx = Arc::new(Mutex::new(Some(vm_tx))); + + let manager_thread = thread::spawn(move || { + manager_event_loop(event_rx, vm_input_tx, 50).expect("event loop"); + }); + + event_tx.send(ManagerEvent::Inc(None)).unwrap(); + assert!(vm_rx.recv_timeout(Duration::from_millis(100)).is_err()); + + event_tx.send(ManagerEvent::Dec(None)).unwrap(); + let msg = vm_rx + .recv_timeout(Duration::from_secs(2)) + .expect("poweroff"); + match msg { + VmInput::Bytes(data) => { + assert_eq!(data, b"systemctl poweroff\n"); + } + _ => panic!("unexpected vm input"), + } + let _ = event_tx.send(ManagerEvent::VmExited(None)); + let _ = manager_thread.join(); + } +} + +struct ManagerOptions { + ensure_signed: bool, + detach: bool, + prepare_vm: bool, +} + +trait VmExecutor { + fn run_vm( + &self, + args: vm::CliArgs, + extra_login_actions: Vec, + extra_shares: Vec, + config: Arc>, + instance_dir: PathBuf, + vm_input_tx: Arc>>>, + ) -> Result<(), Box>; +} + +struct RealVmExecutor; + +impl VmExecutor for RealVmExecutor { + fn run_vm( + &self, + args: vm::CliArgs, + extra_login_actions: Vec, + extra_shares: Vec, + config: Arc>, + instance_dir: PathBuf, + vm_input_tx: Arc>>>, + ) -> Result<(), Box> { + let result = vm::run_with_args_and_extras( + args, + |output_monitor, vm_output_fd, vm_input_fd| { + let io_ctx = spawn_manager_io( + config.clone(), + instance_dir.clone(), + output_monitor, + vm_output_fd, + vm_input_fd, + ); + *vm_input_tx.lock().unwrap() = Some(io_ctx.input_tx.clone()); + io_ctx + }, + extra_login_actions, + extra_shares, + ); + result + } +} + +fn run_manager_with( + project_root: &Path, + args: vm::CliArgs, + auto_shutdown_ms: u64, + executor: &dyn VmExecutor, + options: ManagerOptions, +) -> Result<(), Box> { + if options.ensure_signed { + let _had_skip = env::var("VIBEBOX_SKIP_CODESIGN").ok(); + unsafe { + env::remove_var("VIBEBOX_SKIP_CODESIGN"); + } + vm::ensure_signed(); + unsafe { + env::set_var("VIBEBOX_SKIP_CODESIGN", "1"); + } + } + if options.detach { + detach_from_terminal(); + } + + let project_name = project_root + .file_name() + .ok_or("Project directory has no name")? + .to_string_lossy() + .into_owned(); + let instance_dir = ensure_instance_dir(project_root)?; + if options.prepare_vm { + let _ = ensure_ssh_keypair(&instance_dir)?; + } + + let mut config = load_or_create_instance_config(&instance_dir)?; + if config.vm_ipv4.is_some() { + config.vm_ipv4 = None; + write_instance_config(&instance_dir.join("instance.toml"), &config)?; + } + let config = Arc::new(Mutex::new(config)); + + let ssh_guest_dir = format!("/root/{}", GLOBAL_DIR_NAME); + let extra_shares = vec![DirectoryShare::new( + instance_dir.clone(), + ssh_guest_dir.clone().into(), + true, + )?]; + let extra_login_actions = + build_ssh_login_actions(&config, &project_name, ssh_guest_dir.as_str(), "ssh_key"); + + let socket_path = instance_dir.join(VM_MANAGER_SOCKET_NAME); + if let Ok(stream) = UnixStream::connect(&socket_path) { + drop(stream); + return Ok(()); + } + if socket_path.exists() { + let _ = fs::remove_file(&socket_path); + } + + let listener = UnixListener::bind(&socket_path)?; + let _ = fs::set_permissions(&socket_path, fs::Permissions::from_mode(0o600)); + tracing::info!(path = %socket_path.display(), "vm manager socket bound"); + + let (event_tx, event_rx) = mpsc::channel::(); + let event_tx_accept = event_tx.clone(); + thread::spawn(move || { + for stream in listener.incoming() { + match stream { + Ok(stream) => { + let event_tx_conn = event_tx_accept.clone(); + thread::spawn(move || { + let pid = read_client_pid(&stream); + let _ = event_tx_conn.send(ManagerEvent::Inc(pid)); + wait_for_disconnect(stream); + let _ = event_tx_conn.send(ManagerEvent::Dec(pid)); + }); + } + Err(_) => break, + } + } + }); + + let vm_input_tx: Arc>>> = Arc::new(Mutex::new(None)); + let vm_input_for_loop = vm_input_tx.clone(); + let event_loop_handle = + thread::spawn(move || manager_event_loop(event_rx, vm_input_for_loop, auto_shutdown_ms)); + + tracing::info!("vm manager launching vm"); + let vm_result = executor.run_vm( + args, + extra_login_actions, + extra_shares, + config.clone(), + instance_dir.clone(), + vm_input_tx.clone(), + ); + tracing::info!("vm manager vm run completed"); + let vm_err = vm_result.err().map(|e| e.to_string()); + let _ = event_tx.send(ManagerEvent::VmExited(vm_err.clone())); + let event_loop_result: Result<(), String> = event_loop_handle + .join() + .unwrap_or_else(|_| Err("vm manager event loop panicked".into())) + .map_err(|err| err.to_string()); + let _ = fs::remove_file(&socket_path); + if let Err(err) = &event_loop_result { + tracing::error!(error = %err, "vm manager exiting due to event loop error"); + return Err(err.to_string().into()); + } + if let Some(err) = vm_err { + tracing::error!(error = %err, "vm manager exiting due to vm error"); + return Err(err.into()); + } + tracing::info!("vm manager exiting"); + Ok(event_loop_result?) +} + +fn manager_event_loop( + event_rx: mpsc::Receiver, + vm_input_tx: Arc>>>, + auto_shutdown_ms: u64, +) -> Result<(), String> { + let mut ref_count: usize = 0; + let mut shutdown_deadline: Option = None; + let mut shutdown_sent = false; + let grace = Duration::from_millis(auto_shutdown_ms.max(1)); + + loop { + let timeout = match shutdown_deadline { + Some(deadline) => deadline.saturating_duration_since(Instant::now()), + None => Duration::from_secs(1), + }; + + match event_rx.recv_timeout(timeout) { + Ok(ManagerEvent::Inc(pid)) => { + ref_count = ref_count.saturating_add(1); + tracing::info!( + ref_count, + pid = pid.unwrap_or(0), + pid_known = pid.is_some(), + "vm manager refcount increment" + ); + shutdown_deadline = None; + shutdown_sent = false; + } + Ok(ManagerEvent::Dec(pid)) => { + if ref_count > 0 { + ref_count -= 1; + } + tracing::info!( + ref_count, + pid = pid.unwrap_or(0), + pid_known = pid.is_some(), + "vm manager refcount decrement" + ); + if ref_count == 0 { + shutdown_deadline = Some(Instant::now() + grace); + tracing::info!(grace_ms = auto_shutdown_ms, "shutdown scheduled"); + } + } + Ok(ManagerEvent::VmExited(err)) => { + if let Some(err) = err { + tracing::error!(error = %err, "vm exited with error"); + } + break; + } + Err(mpsc::RecvTimeoutError::Timeout) => { + if let Some(deadline) = shutdown_deadline { + if Instant::now() >= deadline && !shutdown_sent { + if let Some(tx) = vm_input_tx.lock().unwrap().clone() { + let _ = tx.send(VmInput::Bytes(b"systemctl poweroff\n".to_vec())); + } + tracing::info!("shutdown command sent"); + shutdown_sent = true; + shutdown_deadline = None; + } + } + } + Err(mpsc::RecvTimeoutError::Disconnected) => break, + } + } + + Ok(()) +}