diff --git a/Cargo.lock b/Cargo.lock index e935d9dc52..4d2b4b54a6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2076,6 +2076,7 @@ dependencies = [ "google-cloud-auth", "http 1.4.0", "libsqlite3-sys", + "mockito", "oauth2", "pretty_assertions", "reqwest 0.12.28", diff --git a/Cargo.toml b/Cargo.toml index 9a9f43786a..e79b4065ee 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,7 +29,7 @@ aws-smithy-runtime = { version = "1.10", features = ["connector-hyper-0-14-x", " base64 = "0.22.1" bytes = "1.11.1" chrono = { version = "0.4.44", features = ["serde"] } -clap = { version = "4.6.0", features = ["derive"] } +clap = { version = "4.6.0", features = ["derive", "env"] } clap_complete = "4.6.0" colored = "3.1.1" console = "0.16.3" diff --git a/crates/forge_api/src/forge_api.rs b/crates/forge_api/src/forge_api.rs index 1aa7ff6d91..bcdb8f1711 100644 --- a/crates/forge_api/src/forge_api.rs +++ b/crates/forge_api/src/forge_api.rs @@ -11,7 +11,7 @@ use forge_app::{ ProviderAuthService, ProviderService, Services, User, UserUsage, Walker, WorkspaceService, }; use forge_domain::{Agent, ConsoleWriter, *}; -use forge_infra::ForgeInfra; +use forge_infra::{ForgeInfra, TensorlakeConfig}; use forge_repo::ForgeRepo; use forge_services::ForgeServices; use forge_stream::MpscStream; @@ -40,6 +40,7 @@ impl ForgeAPI { } impl ForgeAPI>, ForgeRepo> { + /// Initialises the API with the local command executor (default mode). pub fn init(cwd: PathBuf) -> Self { let infra = Arc::new(ForgeInfra::new(cwd)); let repo = Arc::new(ForgeRepo::new(infra.clone())); @@ -47,6 +48,15 @@ impl ForgeAPI>, ForgeRepo> { ForgeAPI::new(app, repo) } + /// Initialises the API routing all shell commands through an isolated + /// Tensorlake Firecracker microVM sandbox. + pub fn init_with_tensorlake(cwd: PathBuf, config: TensorlakeConfig) -> Self { + let infra = Arc::new(ForgeInfra::new_with_tensorlake(cwd, config)); + let repo = Arc::new(ForgeRepo::new(infra.clone())); + let app = Arc::new(ForgeServices::new(repo.clone())); + ForgeAPI::new(app, repo) + } + pub async fn get_skills_internal(&self) -> Result> { use forge_domain::SkillRepository; self.infra.load_skills().await diff --git a/crates/forge_api/src/lib.rs b/crates/forge_api/src/lib.rs index bf407ba049..5b686cf4cc 100644 --- a/crates/forge_api/src/lib.rs +++ b/crates/forge_api/src/lib.rs @@ -6,3 +6,4 @@ pub use forge_api::*; pub use forge_app::dto::*; pub use forge_app::{Plan, UsageInfo, UserUsage}; pub use forge_domain::{Agent, *}; +pub use forge_infra::TensorlakeConfig; diff --git a/crates/forge_infra/Cargo.toml b/crates/forge_infra/Cargo.toml index 6f3968f3dd..8cfce85cd9 100644 --- a/crates/forge_infra/Cargo.toml +++ b/crates/forge_infra/Cargo.toml @@ -54,3 +54,4 @@ serial_test = "3.4" fake = { version = "5.1.0", features = ["derive"] } pretty_assertions.workspace = true forge_domain = { path = "../forge_domain" } +mockito = { workspace = true } diff --git a/crates/forge_infra/src/forge_infra.rs b/crates/forge_infra/src/forge_infra.rs index b899eee341..663e3f386f 100644 --- a/crates/forge_infra/src/forge_infra.rs +++ b/crates/forge_infra/src/forge_infra.rs @@ -31,8 +31,20 @@ use crate::http::ForgeHttpInfra; use crate::inquire::ForgeInquire; use crate::mcp_client::ForgeMcpClient; use crate::mcp_server::ForgeMcpServer; +use crate::tensorlake::{TensorlakeCommandExecutor, TensorlakeConfig}; use crate::walker::ForgeWalkerService; +/// Abstraction over the available command execution backends. +/// +/// Commands are either executed locally on the host machine via +/// `ForgeCommandExecutorService`, or inside an isolated Tensorlake +/// Firecracker microVM via `TensorlakeCommandExecutor`. +#[derive(Clone)] +enum CommandExecutor { + Local(Arc), + Tensorlake(Arc), +} + #[derive(Clone)] pub struct ForgeInfra { // TODO: Drop the "Service" suffix. Use names like ForgeFileReader, ForgeFileWriter, @@ -44,7 +56,7 @@ pub struct ForgeInfra { file_meta_service: Arc, create_dirs_service: Arc, directory_reader_service: Arc, - command_executor_service: Arc, + command_executor: CommandExecutor, inquire_service: Arc, mcp_server: ForgeMcpServer, walker_service: Arc, @@ -55,6 +67,8 @@ pub struct ForgeInfra { } impl ForgeInfra { + /// Creates a `ForgeInfra` instance that executes shell commands locally on + /// the host machine. pub fn new(cwd: PathBuf) -> Self { let config_infra = Arc::new(ForgeEnvironmentInfra::new(cwd)); let env = config_infra.get_environment(); @@ -76,9 +90,49 @@ impl ForgeInfra { file_meta_service, create_dirs_service: Arc::new(ForgeCreateDirsService), directory_reader_service, - command_executor_service: Arc::new(ForgeCommandExecutorService::new( + command_executor: CommandExecutor::Local(Arc::new(ForgeCommandExecutorService::new( env.clone(), output_printer.clone(), + ))), + inquire_service: Arc::new(ForgeInquire::new()), + mcp_server: ForgeMcpServer, + walker_service: Arc::new(ForgeWalkerService::new()), + strategy_factory: Arc::new(ForgeAuthStrategyFactory::new()), + http_service, + grpc_client, + output_printer, + } + } + + /// Creates a `ForgeInfra` instance that executes shell commands inside an + /// isolated Tensorlake Firecracker microVM sandbox. + /// + /// A single sandbox is provisioned lazily on the first command execution + /// and reused for the entire session. The sandbox is terminated when + /// the returned `ForgeInfra` is dropped. + pub fn new_with_tensorlake(cwd: PathBuf, config: TensorlakeConfig) -> Self { + let config_infra = Arc::new(ForgeEnvironmentInfra::new(cwd)); + let env = config_infra.get_environment(); + + let file_write_service = Arc::new(ForgeFileWriteService::new()); + let http_service = Arc::new(ForgeHttpInfra::new(env.clone(), file_write_service.clone())); + let file_read_service = Arc::new(ForgeFileReadService::new()); + let file_meta_service = Arc::new(ForgeFileMetaService); + let directory_reader_service = + Arc::new(ForgeDirectoryReaderService::new(env.parallel_file_reads)); + let grpc_client = Arc::new(ForgeGrpcClient::new(env.service_url.clone())); + let output_printer = Arc::new(StdConsoleWriter::default()); + + Self { + file_read_service, + file_write_service, + file_remove_service: Arc::new(ForgeFileRemoveService::new()), + config_infra, + file_meta_service, + create_dirs_service: Arc::new(ForgeCreateDirsService), + directory_reader_service, + command_executor: CommandExecutor::Tensorlake(Arc::new( + TensorlakeCommandExecutor::new(config), )), inquire_service: Arc::new(ForgeInquire::new()), mcp_server: ForgeMcpServer, @@ -196,9 +250,16 @@ impl CommandInfra for ForgeInfra { silent: bool, env_vars: Option>, ) -> anyhow::Result { - self.command_executor_service - .execute_command(command, working_dir, silent, env_vars) - .await + match &self.command_executor { + CommandExecutor::Local(svc) => { + svc.execute_command(command, working_dir, silent, env_vars) + .await + } + CommandExecutor::Tensorlake(svc) => { + svc.execute_command(command, working_dir, silent, env_vars) + .await + } + } } async fn execute_command_raw( @@ -207,9 +268,16 @@ impl CommandInfra for ForgeInfra { working_dir: PathBuf, env_vars: Option>, ) -> anyhow::Result { - self.command_executor_service - .execute_command_raw(command, working_dir, env_vars) - .await + match &self.command_executor { + CommandExecutor::Local(svc) => { + svc.execute_command_raw(command, working_dir, env_vars) + .await + } + CommandExecutor::Tensorlake(svc) => { + svc.execute_command_raw(command, working_dir, env_vars) + .await + } + } } } diff --git a/crates/forge_infra/src/lib.rs b/crates/forge_infra/src/lib.rs index 542c0f0891..b3d503fad8 100644 --- a/crates/forge_infra/src/lib.rs +++ b/crates/forge_infra/src/lib.rs @@ -1,6 +1,7 @@ mod console; mod env; pub mod executor; +pub mod tensorlake; mod auth; mod error; @@ -24,3 +25,4 @@ pub use env::ForgeEnvironmentInfra; pub use executor::ForgeCommandExecutorService; pub use forge_infra::*; pub use kv_storage::CacacheStorage; +pub use tensorlake::{TensorlakeCommandExecutor, TensorlakeConfig}; diff --git a/crates/forge_infra/src/tensorlake.rs b/crates/forge_infra/src/tensorlake.rs new file mode 100644 index 0000000000..44f31c11a5 --- /dev/null +++ b/crates/forge_infra/src/tensorlake.rs @@ -0,0 +1,564 @@ +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Arc; + +use anyhow::{Context, anyhow}; +use async_trait::async_trait; +use forge_app::CommandInfra; +use forge_domain::CommandOutput; +use serde::{Deserialize, Serialize}; +use tokio::sync::Mutex; + +const TENSORLAKE_API_BASE: &str = "https://api.tensorlake.ai"; + +/// Configuration for a Tensorlake sandbox session. +#[derive(Debug, Clone)] +pub struct TensorlakeConfig { + /// Tensorlake API key used for all requests. + pub api_key: String, + /// Number of vCPUs to allocate for the sandbox (default: 2.0). + pub cpus: f64, + /// Memory in megabytes to allocate for the sandbox (default: 4096). + pub memory_mb: u64, + /// Inactivity timeout in seconds before the sandbox auto-suspends (default: + /// 3600). + pub timeout_secs: u64, + /// Base URL for the Tensorlake API (default: `https://api.tensorlake.ai`). + /// Overridable in tests to point at a local mock server. + pub base_url: String, +} + +impl TensorlakeConfig { + /// Creates a new `TensorlakeConfig` with the given API key and sensible + /// defaults. + pub fn new(api_key: String) -> Self { + Self { + api_key, + cpus: 2.0, + memory_mb: 4096, + timeout_secs: 3600, + base_url: TENSORLAKE_API_BASE.to_string(), + } + } +} + +/// Response returned by the Tensorlake sandboxes create endpoint. +#[derive(Debug, Deserialize)] +struct CreateSandboxResponse { + sandbox_id: String, +} + +/// Response returned by the per-sandbox process execution endpoint. +#[derive(Debug, Deserialize)] +struct StartProcessResponse { + pid: u64, +} + +/// Combined output response from the per-sandbox output endpoint. +#[derive(Debug, Deserialize)] +struct ProcessOutputResponse { + lines: Vec, +} + +/// Request body for starting a process inside a sandbox. +#[derive(Debug, Serialize)] +struct StartProcessRequest<'a> { + command: &'a str, + args: Vec<&'a str>, + working_dir: String, + /// Optional environment variables to inject, as `KEY=VALUE` pairs parsed + /// into a map. The Tensorlake API accepts `{"KEY": "VALUE"}` format. + #[serde(skip_serializing_if = "Option::is_none")] + env: Option>, +} + +/// Owns the sandbox lifetime and issues the DELETE on drop. +/// +/// Wrapped in `Arc` inside `TensorlakeCommandExecutor` so that the sandbox is +/// terminated exactly once — when the last clone of the executor is dropped and +/// the `Arc` reference count reaches zero. +struct SandboxGuard { + sandbox_id: Arc>>, + client: reqwest::Client, + api_key: String, + base_url: String, +} + +impl Drop for SandboxGuard { + /// Synchronously terminates the sandbox by sending a DELETE request. + /// + /// Uses `block_in_place` + `block_on` so the HTTP call completes before + /// `Drop` returns, ensuring the sandbox is always cleaned up on normal + /// exit, Ctrl-C, and panics that unwind the stack. SIGKILL remains + /// unhandled by design; `timeout_secs` acts as the safety net there. + fn drop(&mut self) { + let sandbox_id = self.sandbox_id.clone(); + let client = self.client.clone(); + let api_key = self.api_key.clone(); + let base_url = self.base_url.clone(); + + if let Ok(handle) = tokio::runtime::Handle::try_current() { + tokio::task::block_in_place(|| { + handle.block_on(async move { + let guard = sandbox_id.lock().await; + if let Some(id) = guard.as_deref() { + let url = format!("{}/sandboxes/{}", base_url, id); + let _ = client.delete(&url).bearer_auth(&api_key).send().await; + tracing::debug!(sandbox_id = %id, "Tensorlake sandbox terminated"); + } + }); + }); + } + } +} + +/// Infrastructure implementation that executes shell commands inside an +/// isolated Tensorlake Firecracker microVM sandbox. +/// +/// A single sandbox is created lazily on the first command execution and reused +/// for the lifetime of the `TensorlakeCommandExecutor` instance. The sandbox is +/// terminated when the last clone of the executor is dropped. +#[derive(Clone)] +pub struct TensorlakeCommandExecutor { + config: TensorlakeConfig, + client: reqwest::Client, + /// Lazily initialized sandbox ID, shared across clones via `Arc>`. + sandbox_id: Arc>>, + /// Dropping the last clone of this `Arc` triggers `SandboxGuard::drop`, + /// which issues the sandbox DELETE exactly once. + _guard: Arc, +} + +impl TensorlakeCommandExecutor { + /// Creates a new executor with the provided Tensorlake configuration. + pub fn new(config: TensorlakeConfig) -> Self { + let client = reqwest::Client::new(); + let sandbox_id = Arc::new(Mutex::new(None)); + let guard = Arc::new(SandboxGuard { + sandbox_id: sandbox_id.clone(), + client: client.clone(), + api_key: config.api_key.clone(), + base_url: config.base_url.clone(), + }); + Self { config, client, sandbox_id, _guard: guard } + } + + /// Returns the sandbox ID, creating a new sandbox if one has not been + /// provisioned yet for this session. + async fn ensure_sandbox(&self) -> anyhow::Result { + let mut guard = self.sandbox_id.lock().await; + if let Some(id) = guard.as_deref() { + return Ok(id.to_string()); + } + + let id = self.create_sandbox().await?; + tracing::info!(sandbox_id = %id, "Tensorlake sandbox created"); + *guard = Some(id.clone()); + Ok(id) + } + + /// Provisions a new Tensorlake sandbox and returns its ID. + async fn create_sandbox(&self) -> anyhow::Result { + let url = format!("{}/sandboxes", self.config.base_url); + let body = serde_json::json!({ + "resources": { + "cpus": self.config.cpus, + "memory_mb": self.config.memory_mb, + }, + "timeout_secs": self.config.timeout_secs, + }); + + let response = self + .client + .post(&url) + .bearer_auth(&self.config.api_key) + .json(&body) + .send() + .await + .context("Failed to send create sandbox request to Tensorlake")?; + + if !response.status().is_success() { + let status = response.status(); + let text = response.text().await.unwrap_or_default(); + return Err(anyhow!( + "Tensorlake sandbox creation failed with status {status}: {text}" + )); + } + + let parsed: CreateSandboxResponse = response + .json() + .await + .context("Failed to parse Tensorlake create sandbox response")?; + + // Poll until the sandbox is running (it starts as "pending") + self.wait_for_running(&parsed.sandbox_id).await?; + + Ok(parsed.sandbox_id) + } + + /// Polls the sandbox status endpoint until the sandbox reaches the + /// "running" state. + async fn wait_for_running(&self, sandbox_id: &str) -> anyhow::Result<()> { + let url = format!("{}/sandboxes/{}", self.config.base_url, sandbox_id); + for attempt in 0..60 { + tokio::time::sleep(std::time::Duration::from_secs(if attempt == 0 { + 1 + } else { + 2 + })) + .await; + + let resp = self + .client + .get(&url) + .bearer_auth(&self.config.api_key) + .send() + .await + .context("Failed to poll sandbox status")?; + + if !resp.status().is_success() { + continue; + } + + let body: serde_json::Value = resp + .json() + .await + .context("Failed to parse sandbox status")?; + + match body.get("status").and_then(|s| s.as_str()) { + Some("running") => return Ok(()), + Some("terminated") => { + return Err(anyhow!("Tensorlake sandbox terminated unexpectedly")); + } + _ => continue, + } + } + Err(anyhow!( + "Timed out waiting for Tensorlake sandbox to become running" + )) + } + + /// Returns the per-sandbox proxy base URL for API calls. + fn sandbox_proxy_url(&self, sandbox_id: &str) -> String { + format!("https://{}.sandbox.tensorlake.ai/api/v1", sandbox_id) + } +} + +#[async_trait] +impl CommandInfra for TensorlakeCommandExecutor { + /// Executes a shell command inside the Tensorlake sandbox and returns the + /// captured output. + async fn execute_command( + &self, + command: String, + working_dir: PathBuf, + _silent: bool, + env_vars: Option>, + ) -> anyhow::Result { + let sandbox_id = self.ensure_sandbox().await?; + let proxy = self.sandbox_proxy_url(&sandbox_id); + + // The sandbox is a remote Linux microVM. Host-specific paths (e.g. macOS + // `/Users/…`) do not exist inside the sandbox. Fall back to `/tmp` so + // that process spawn never fails with "No such file or directory". + + let cwd = { + let host_path = working_dir.to_string_lossy(); + // Handle both Unix and Windows paths + if host_path.starts_with("/Users/") || host_path.starts_with("/home/") || host_path.contains(":\\") { + "/tmp".to_string() + } else { + host_path.into_owned() + } + }; + + + // Parse `KEY=VALUE` strings into the dict format the Tensorlake API expects. + let env = env_vars.map(|vars| { + vars.into_iter() + .filter_map(|kv| { + let mut parts = kv.splitn(2, '='); + let key = parts.next()?.to_string(); + let value = parts.next().unwrap_or("").to_string(); + Some((key, value)) + }) + .collect::>() + }); + + // Use `sh -c ` so pipes and redirects work correctly. + let start_url = format!("{}/processes", proxy); + let request = StartProcessRequest { + command: "sh", + args: vec!["-c", &command], + working_dir: cwd, + env, + }; + + tracing::info!(command = %command, sandbox_id = %sandbox_id, "Executing command in Tensorlake sandbox"); + + let response = self + .client + .post(&start_url) + .bearer_auth(&self.config.api_key) + .json(&request) + .send() + .await + .context("Failed to start process in Tensorlake sandbox")?; + + if !response.status().is_success() { + let status = response.status(); + let text = response.text().await.unwrap_or_default(); + return Err(anyhow!( + "Tensorlake process start failed with status {status}: {text}" + )); + } + + let started: StartProcessResponse = response + .json() + .await + .context("Failed to parse Tensorlake process start response")?; + + let pid = started.pid; + + // Poll until the process exits. + let exit_code = self.wait_for_process_exit(&proxy, pid).await?; + + // Collect stdout and stderr. + let stdout = self.get_process_output(&proxy, pid, "stdout").await?; + let stderr = self.get_process_output(&proxy, pid, "stderr").await?; + + Ok(CommandOutput { stdout, stderr, exit_code: Some(exit_code), command }) + } + + /// Interactive (raw) commands are not supported in Tensorlake sandbox mode. + /// + /// Raw command execution requires an attached TTY which is not available + /// over the Tensorlake HTTP API. This method always returns an error + /// directing the caller to use `execute_command` instead. + async fn execute_command_raw( + &self, + _command: &str, + _working_dir: PathBuf, + _env_vars: Option>, + ) -> anyhow::Result { + Err(anyhow!( + "Interactive (raw) command execution is not supported in Tensorlake sandbox mode. \ + Use non-interactive commands instead." + )) + } +} + +impl TensorlakeCommandExecutor { + /// Polls the process status until it exits, returning the exit code. + async fn wait_for_process_exit(&self, proxy: &str, pid: u64) -> anyhow::Result { + let url = format!("{}/processes/{}", proxy, pid); + for _ in 0..300 { + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + + let resp = self + .client + .get(&url) + .bearer_auth(&self.config.api_key) + .send() + .await + .context("Failed to poll process status")?; + + if !resp.status().is_success() { + let status = resp.status(); + return Err(anyhow!( + "Failed to poll process {pid} status: HTTP {status}" + )); + } + + let body: serde_json::Value = resp + .json() + .await + .context("Failed to parse process status")?; + + if body.get("status").and_then(|s| s.as_str()) == Some("exited") { + let code = body.get("exit_code").and_then(|c| c.as_i64()).unwrap_or(0) as i32; + return Ok(code); + } + } + Err(anyhow!( + "Timed out waiting for Tensorlake process {} to exit", + pid + )) + } + + /// Fetches the captured output lines (stdout or stderr) for a completed + /// process. + async fn get_process_output( + &self, + proxy: &str, + pid: u64, + stream: &str, + ) -> anyhow::Result { + let url = format!("{}/processes/{}/{}", proxy, pid, stream); + let resp = self + .client + .get(&url) + .bearer_auth(&self.config.api_key) + .send() + .await + .with_context(|| format!("Failed to fetch process {stream} for pid {pid}"))?; + + if !resp.status().is_success() { + let status = resp.status(); + return Err(anyhow!( + "Failed to fetch {stream} for pid {pid}: HTTP {status}" + )); + } + + let output: ProcessOutputResponse = resp + .json() + .await + .with_context(|| format!("Failed to parse process {stream} response"))?; + + Ok(output.lines.join("\n")) + } +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn test_tensorlake_config_defaults() { + let fixture = TensorlakeConfig::new("test-api-key".to_string()); + + assert_eq!(fixture.api_key, "test-api-key"); + assert_eq!(fixture.cpus, 2.0); + assert_eq!(fixture.memory_mb, 4096); + assert_eq!(fixture.timeout_secs, 3600); + assert_eq!(fixture.base_url, "https://api.tensorlake.ai"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_tensorlake_executor_creation() { + let config = TensorlakeConfig::new("test-api-key".to_string()); + let executor = TensorlakeCommandExecutor::new(config.clone()); + + assert_eq!(executor.config.api_key, config.api_key); + assert_eq!(executor.config.cpus, config.cpus); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_sandbox_proxy_url() { + let config = TensorlakeConfig::new("key".to_string()); + let executor = TensorlakeCommandExecutor::new(config); + let url = executor.sandbox_proxy_url("abc123"); + assert_eq!(url, "https://abc123.sandbox.tensorlake.ai/api/v1"); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_cleanup_fires_on_last_clone_dropped() { + let config = TensorlakeConfig::new("key".to_string()); + let executor = TensorlakeCommandExecutor::new(config); + let clone = executor.clone(); + + // Both the original and the clone share the same Arc. + assert_eq!(Arc::strong_count(&executor._guard), 2); + + // Dropping the clone reduces the count to 1 — cleanup not yet triggered. + drop(clone); + assert_eq!(Arc::strong_count(&executor._guard), 1); + + // Dropping the original reduces the count to 0 — cleanup fires here. + // (No sandbox was provisioned so the spawned task is a no-op.) + drop(executor); + } + + /// Verifies that dropping the executor sends a DELETE /sandboxes/{id} + /// request synchronously — i.e. the request completes before `drop` + /// returns, so the mock server receives it before `assert_hits` is + /// checked. + #[tokio::test(flavor = "multi_thread")] + async fn test_sandbox_terminated_on_drop() { + let mut server = mockito::Server::new_async().await; + + let mock = server + .mock("DELETE", "/sandboxes/test-sandbox-id") + .with_status(200) + .expect(1) + .create_async() + .await; + + let mut config = TensorlakeConfig::new("test-key".to_string()); + config.base_url = server.url(); + + let executor = TensorlakeCommandExecutor::new(config); + + // Manually plant a sandbox ID so the guard has something to DELETE. + { + let mut guard = executor.sandbox_id.lock().await; + *guard = Some("test-sandbox-id".to_string()); + } + + // Drop the executor — SandboxGuard::drop must complete the DELETE before + // returning, so the mock expectation is satisfied synchronously. + drop(executor); + + mock.assert_async().await; + } + + /// Verifies that a clone and the original share the guard, and the DELETE + /// is sent exactly once — only when the last clone is dropped. + #[tokio::test(flavor = "multi_thread")] + async fn test_sandbox_terminated_exactly_once_across_clones() { + let mut server = mockito::Server::new_async().await; + + let mock = server + .mock("DELETE", "/sandboxes/shared-sandbox") + .with_status(200) + .expect(1) // must fire exactly once, not twice + .create_async() + .await; + + let mut config = TensorlakeConfig::new("test-key".to_string()); + config.base_url = server.url(); + + let executor = TensorlakeCommandExecutor::new(config); + let clone = executor.clone(); + + { + let mut guard = executor.sandbox_id.lock().await; + *guard = Some("shared-sandbox".to_string()); + } + + // Dropping the clone must NOT trigger the DELETE yet. + drop(clone); + // Dropping the original must trigger the DELETE exactly once. + drop(executor); + + mock.assert_async().await; + } + + #[test] + fn test_env_vars_parsed_into_map() { + let vars = vec![ + "FOO=bar".to_string(), + "BAZ=qux=with=equals".to_string(), + "EMPTY=".to_string(), + ]; + + let actual: HashMap = vars + .into_iter() + .filter_map(|kv| { + let mut parts = kv.splitn(2, '='); + let key = parts.next()?.to_string(); + let value = parts.next().unwrap_or("").to_string(); + Some((key, value)) + }) + .collect(); + + let mut expected = HashMap::new(); + expected.insert("FOO".to_string(), "bar".to_string()); + expected.insert("BAZ".to_string(), "qux=with=equals".to_string()); + expected.insert("EMPTY".to_string(), "".to_string()); + + assert_eq!(actual, expected); + } +} diff --git a/crates/forge_main/src/cli.rs b/crates/forge_main/src/cli.rs index d2189594f6..f160ae796f 100644 --- a/crates/forge_main/src/cli.rs +++ b/crates/forge_main/src/cli.rs @@ -50,6 +50,15 @@ pub struct Cli { #[arg(long)] pub sandbox: Option, + /// Run all shell commands inside an isolated Tensorlake Firecracker microVM + /// sandbox. + /// + /// Requires a valid Tensorlake API key. Can also be supplied via the + /// TENSORLAKE_API_KEY environment variable. When set, every shell tool call + /// is routed through a remote Tensorlake sandbox instead of the local host. + #[arg(long, env = "TENSORLAKE_API_KEY")] + pub tensorlake: Option, + /// Enable verbose logging output. #[arg(long, default_value_t = false)] pub verbose: bool, diff --git a/crates/forge_main/src/main.rs b/crates/forge_main/src/main.rs index 3f680bdb4d..2cd4a4db23 100644 --- a/crates/forge_main/src/main.rs +++ b/crates/forge_main/src/main.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use anyhow::Result; use clap::Parser; -use forge_api::ForgeAPI; +use forge_api::{ForgeAPI, TensorlakeConfig}; use forge_domain::TitleFormat; use forge_main::{Cli, Sandbox, TitleDisplayExt, UI, tracker}; @@ -19,7 +19,11 @@ async fn main() -> Result<()> { // available let _ = rustls::crypto::ring::default_provider().install_default(); - // Set up panic hook for better error display + // Set up panic hook for better error display. + // Important: do NOT call `std::process::exit` here. Exiting inside the hook + // skips stack unwinding, which prevents `Drop` implementations (such as the + // Tensorlake `SandboxGuard`) from running. Returning from the hook lets Rust + // unwind the stack normally so all destructors fire before the process exits. panic::set_hook(Box::new(|panic_info| { let message = if let Some(s) = panic_info.payload().downcast_ref::<&str>() { s.to_string() @@ -31,7 +35,6 @@ async fn main() -> Result<()> { println!("{}", TitleFormat::error(message.to_string()).display()); tracker::error_blocking(message); - std::process::exit(1); })); // Initialize and run the UI @@ -62,7 +65,17 @@ async fn main() -> Result<()> { (_, _) => std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")), }; - let mut ui = UI::init(cli, move || ForgeAPI::init(cwd.clone()))?; + let tensorlake_key = cli.tensorlake.clone(); + let mut ui = UI::init(cli, move || { + if let Some(api_key) = &tensorlake_key { + ForgeAPI::init_with_tensorlake( + cwd.clone(), + TensorlakeConfig::new(api_key.to_string()), + ) + } else { + ForgeAPI::init(cwd.clone()) + } + })?; ui.run().await; Ok(())