implement persistent worker mode in CLI and SDK to reduce subprocess

crates/capsule-cli/Cargo.toml

@@ -17,6 +17,7 @@ path = "src/main.rs"
 [dependencies]
 clap = { version = "4.5.53", features = ["derive"] }
 capsule-core = { version= "0.8.8", path = "../capsule-core" }
-tokio = { version = "1.48.0", features = ["rt", "rt-multi-thread", "macros"] }
+tokio = { version = "1.48.0", features = ["rt", "rt-multi-thread", "macros", "io-util", "sync"] }
+serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 dotenvy = "0.15.7"

crates/capsule-cli/src/cli.rs

@@ -35,6 +35,7 @@ pub enum Commands {
         #[arg(long)]
         export: bool,
     },
+    Worker,
     Exec {
         file: String,
 

crates/capsule-cli/src/commands/exec.rs

@@ -1,5 +1,6 @@
 use std::fmt;
 use std::path::Path;
+use std::sync::Arc;
 use std::time::Instant;
 
 use capsule_core::config::manifest::{Manifest, ManifestError};
@@ -49,6 +50,17 @@ pub async fn execute(
     mounts: Vec<String>,
     json: bool,
     verbose: bool,
+) -> Result<String, ExecError> {
+    execute_with_runtime(wasm_path, args, mounts, json, verbose, None).await
+}
+
+pub async fn execute_with_runtime(
+    wasm_path: &Path,
+    args: Vec<String>,
+    mounts: Vec<String>,
+    json: bool,
+    verbose: bool,
+    shared_runtime: Option<Arc<Runtime>>,
 ) -> Result<String, ExecError> {
     let ext = wasm_path.extension().and_then(|e| e.to_str()).unwrap_or("");
 
@@ -82,24 +94,27 @@ pub async fn execute(
 
     load_env_variables(&project_root).map_err(ExecError::IoError)?;
 
-    reporter.start_progress("Initializing runtime");
-
-    let capsule_toml = Manifest::new().map(|m| m.capsule_toml).unwrap_or_default();
-
-    let runtime_config = RuntimeConfig { cache_dir, verbose };
-
     let mut execution_policy = ExecutionPolicy::default().compute(Some(Compute::Custom(u64::MAX)));
 
     execution_policy.mounts.extend(mounts);
 
-    let runtime = Runtime::new(runtime_config, capsule_toml)?;
+    let runtime = match shared_runtime {
+        Some(r) => r,
+        None => {
+            reporter.start_progress("Initializing runtime");
+            let capsule_toml = Manifest::new().map(|m| m.capsule_toml).unwrap_or_default();
+            let runtime_config = RuntimeConfig { cache_dir, verbose };
+            let r = Runtime::new(runtime_config, capsule_toml)?;
+            reporter.finish_progress(Some("Runtime ready"));
+            r
+        }
+    };
 
     let create_instance_command = CreateInstance::new(execution_policy.clone(), args.clone())
         .wasm_path(wasm_path_abs)
         .project_root(project_root);
 
     let (store, instance, task_id) = runtime.execute(create_instance_command).await?;
-    reporter.finish_progress(Some("Runtime ready"));
 
     let start_time = Instant::now();
 

crates/capsule-cli/src/commands/mod.rs

@@ -2,6 +2,7 @@ pub mod build;
 pub mod exec;
 pub mod run;
 pub mod shared;
+pub mod worker;
 
 pub use build::BuildError;
 pub use exec::ExecError;

crates/capsule-cli/src/commands/run.rs

@@ -1,5 +1,6 @@
 use std::fmt;
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 use std::time::Instant;
 
 use capsule_core::config::manifest::{CapsuleToml, Manifest, ManifestError};
@@ -68,6 +69,17 @@ pub async fn execute(
     mounts: Vec<String>,
     json: bool,
     verbose: bool,
+) -> Result<String, RunError> {
+    execute_with_runtime(file_path, args, mounts, json, verbose, None).await
+}
+
+pub async fn execute_with_runtime(
+    file_path: Option<&Path>,
+    args: Vec<String>,
+    mounts: Vec<String>,
+    json: bool,
+    verbose: bool,
+    shared_runtime: Option<Arc<Runtime>>,
 ) -> Result<String, RunError> {
     let manifest = Manifest::new()?;
     let mut reporter = TaskReporter::new(!json);
@@ -81,8 +93,6 @@ pub async fn execute(
     let compile_result = compile_to_wasm(&file_path, false)?;
     reporter.finish_progress(Some("Environment ready"));
 
-    reporter.start_progress("Initializing runtime");
-
     let project_root = std::env::current_dir().unwrap_or_else(|_| {
         file_path
             .canonicalize()
@@ -93,25 +103,29 @@ pub async fn execute(
 
     load_env_variables(&project_root).map_err(RunError::IoError)?;
 
-    let runtime_config = RuntimeConfig {
-        cache_dir: compile_result.cache_dir,
-        verbose,
-    };
-
     let mut execution_policy =
         extract_main_execution_policy(compile_result.task_registry, &manifest.capsule_toml)
             .unwrap_or_else(|| ExecutionPolicy::default().compute(Some(Compute::Custom(u64::MAX))));
 
     execution_policy.mounts.extend(mounts);
 
-    let runtime = Runtime::new(runtime_config, manifest.capsule_toml)?;
+    let runtime = match shared_runtime {
+        Some(r) => r,
+        None => {
+            reporter.start_progress("Initializing runtime");
+            let runtime_config = RuntimeConfig {
+                cache_dir: compile_result.cache_dir,
+                verbose,
+            };
+            Runtime::new(runtime_config, manifest.capsule_toml)?
+        }
+    };
 
     let create_instance_command = CreateInstance::new(execution_policy.clone(), args.clone())
         .wasm_path(compile_result.wasm_path)
         .project_root(project_root);
 
     let (store, instance, task_id) = runtime.execute(create_instance_command).await?;
-    reporter.finish_progress(Some("Runtime ready"));
 
     let start_time = Instant::now();
 

crates/capsule-cli/src/commands/worker.rs

@@ -0,0 +1,175 @@
+use std::fmt;
+use std::path::Path;
+use std::sync::Arc;
+
+use capsule_core::config::manifest::Manifest;
+use capsule_core::wasm::runtime::{Runtime, RuntimeConfig};
+use serde::{Deserialize, Serialize};
+use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
+use tokio::sync::mpsc;
+
+use crate::commands::exec::ExecError;
+use crate::commands::run::RunError;
+
+const RESPONSE_CHANNEL_SIZE: usize = 64;
+
+#[derive(Debug)]
+pub enum WorkerError {
+    IoError(String),
+    RuntimeError(String),
+}
+
+impl fmt::Display for WorkerError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            WorkerError::IoError(msg) => write!(f, "{}", msg),
+            WorkerError::RuntimeError(msg) => write!(f, "{}", msg),
+        }
+    }
+}
+
+impl From<std::io::Error> for WorkerError {
+    fn from(err: std::io::Error) -> Self {
+        WorkerError::IoError(err.to_string())
+    }
+}
+
+#[derive(Debug, Deserialize)]
+pub struct WorkerRequest {
+    pub id: String,
+    pub file: String,
+    #[serde(default)]
+    pub args: Vec<String>,
+    #[serde(default)]
+    pub mounts: Vec<String>,
+}
+
+#[derive(Debug, Serialize)]
+pub struct WorkerResponse {
+    pub id: String,
+    #[serde(flatten)]
+    pub result: WorkerResult,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(untagged)]
+pub enum WorkerResult {
+    Ok { output: serde_json::Value },
+    Err { error: String },
+}
+
+pub async fn execute() -> Result<(), WorkerError> {
+    let project_root = std::env::current_dir().unwrap_or_default();
+    let cache_dir = project_root.join(".capsule");
+    std::fs::create_dir_all(&cache_dir)?;
+
+    let capsule_toml = Manifest::new().map(|m| m.capsule_toml).unwrap_or_default();
+    let runtime_config = RuntimeConfig {
+        cache_dir,
+        verbose: false,
+    };
+
+    let runtime = Runtime::new(runtime_config, capsule_toml)
+        .map_err(|e| WorkerError::RuntimeError(e.to_string()))?;
+
+    let stdin = tokio::io::stdin();
+    let mut lines = BufReader::new(stdin).lines();
+
+    let (tx, mut rx) = mpsc::channel::<WorkerResponse>(RESPONSE_CHANNEL_SIZE);
+
+    tokio::spawn(async move {
+        let mut stdout = tokio::io::stdout();
+        while let Some(response) = rx.recv().await {
+            if let Ok(mut json) = serde_json::to_string(&response) {
+                json.push('\n');
+                let _ = stdout.write_all(json.as_bytes()).await;
+                let _ = stdout.flush().await;
+            }
+        }
+    });
+
+    while let Ok(Some(line)) = lines.next_line().await {
+        let line = line.trim().to_string();
+        if line.is_empty() {
+            continue;
+        }
+
+        let tx = tx.clone();
+
+        let request: WorkerRequest = match serde_json::from_str(&line) {
+            Ok(r) => r,
+            Err(e) => {
+                let _ = tx
+                    .send(WorkerResponse {
+                        id: "unknown".to_string(),
+                        result: WorkerResult::Err {
+                            error: format!("Failed to parse request: {}", e),
+                        },
+                    })
+                    .await;
+                continue;
+            }
+        };
+
+        let runtime = Arc::clone(&runtime);
+
+        tokio::spawn(async move {
+            let result = dispatch(request.file, request.args, request.mounts, runtime).await;
+
+            let response = WorkerResponse {
+                id: request.id,
+                result: match result {
+                    Ok(output) => match serde_json::from_str(&output) {
+                        Ok(value) => WorkerResult::Ok { output: value },
+                        Err(_) => WorkerResult::Ok {
+                            output: serde_json::Value::String(output),
+                        },
+                    },
+                    Err(e) => WorkerResult::Err { error: e },
+                },
+            };
+
+            let _ = tx.send(response).await;
+        });
+    }
+
+    Ok(())
+}
+
+async fn dispatch(
+    file: String,
+    args: Vec<String>,
+    mounts: Vec<String>,
+    runtime: Arc<Runtime>,
+) -> Result<String, String> {
+    use std::ffi::OsStr;
+
+    let ext = Path::new(&file)
+        .extension()
+        .and_then(OsStr::to_str)
+        .unwrap_or("")
+        .to_lowercase();
+
+    match ext.as_str() {
+        "wasm" | "cwasm" => super::exec::execute_with_runtime(
+            Path::new(&file),
+            args,
+            mounts,
+            true,
+            false,
+            Some(runtime),
+        )
+        .await
+        .map_err(|e: ExecError| e.to_string()),
+        _ => super::run::execute_with_runtime(
+            Some(Path::new(&file)),
+            args,
+            mounts,
+            true,
+            false,
+            Some(runtime),
+        )
+        .await
+        .map_err(|e: RunError| e.to_string()),
+    }
+}

crates/capsule-cli/src/main.rs

@@ -7,7 +7,7 @@ use std::path::Path;
 
 use cli::{Cli, Commands};
 use commands::shared::load_args_file;
-use commands::{BuildError, ExecError, RunError, build, exec, run};
+use commands::{BuildError, ExecError, RunError, build, exec, run, worker};
 
 #[derive(Debug)]
 pub enum CliError {
@@ -72,6 +72,11 @@ async fn main() -> Result<(), CliError> {
                 println!("{}", result);
             }
         }
+        Commands::Worker => {
+            worker::execute()
+                .await
+                .map_err(|e| CliError::RunError(e.to_string()))?;
+        }
         Commands::Build { file, export } => {
             let file_path = file.as_deref().map(Path::new);
             build::execute(file_path, export).await?;