From 5ee2b6d9ba6629bebb1e1dc4f5b9d5f6b5f4ccdb Mon Sep 17 00:00:00 2001 From: Mavdol Date: Fri, 17 Apr 2026 13:49:32 +0200 Subject: [PATCH 1/2] implement persistent worker mode in CLI and SDK to reduce subprocess startup overhead --- crates/capsule-cli/Cargo.lock | 1 + crates/capsule-cli/Cargo.toml | 3 +- crates/capsule-cli/src/cli.rs | 1 + crates/capsule-cli/src/commands/exec.rs | 31 +++- crates/capsule-cli/src/commands/mod.rs | 1 + crates/capsule-cli/src/commands/run.rs | 32 ++-- crates/capsule-cli/src/commands/worker.rs | 175 ++++++++++++++++++++++ crates/capsule-cli/src/main.rs | 7 +- crates/capsule-sdk/javascript/src/run.ts | 162 ++++++++++++++++---- 9 files changed, 364 insertions(+), 49 deletions(-) create mode 100644 crates/capsule-cli/src/commands/worker.rs diff --git a/crates/capsule-cli/Cargo.lock b/crates/capsule-cli/Cargo.lock index ecc00458..62138250 100644 --- a/crates/capsule-cli/Cargo.lock +++ b/crates/capsule-cli/Cargo.lock @@ -356,6 +356,7 @@ dependencies = [ "capsule-core", "clap", "dotenvy", + "serde", "serde_json", "tokio", ] diff --git a/crates/capsule-cli/Cargo.toml b/crates/capsule-cli/Cargo.toml index b80be42b..14d3a9e0 100644 --- a/crates/capsule-cli/Cargo.toml +++ b/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" diff --git a/crates/capsule-cli/src/cli.rs b/crates/capsule-cli/src/cli.rs index 1bb5bffa..368a9236 100644 --- a/crates/capsule-cli/src/cli.rs +++ b/crates/capsule-cli/src/cli.rs @@ -35,6 +35,7 @@ pub enum Commands { #[arg(long)] export: bool, }, + Worker, Exec { file: String, diff --git a/crates/capsule-cli/src/commands/exec.rs b/crates/capsule-cli/src/commands/exec.rs index 611bc0bd..e8625f68 100644 --- a/crates/capsule-cli/src/commands/exec.rs +++ b/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, json: bool, verbose: bool, +) -> Result { + execute_with_runtime(wasm_path, args, mounts, json, verbose, None).await +} + +pub async fn execute_with_runtime( + wasm_path: &Path, + args: Vec, + mounts: Vec, + json: bool, + verbose: bool, + shared_runtime: Option>, ) -> Result { 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(); diff --git a/crates/capsule-cli/src/commands/mod.rs b/crates/capsule-cli/src/commands/mod.rs index 51c82c39..a6f7f013 100644 --- a/crates/capsule-cli/src/commands/mod.rs +++ b/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; diff --git a/crates/capsule-cli/src/commands/run.rs b/crates/capsule-cli/src/commands/run.rs index 1f5edab5..e8233e3b 100644 --- a/crates/capsule-cli/src/commands/run.rs +++ b/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, json: bool, verbose: bool, +) -> Result { + execute_with_runtime(file_path, args, mounts, json, verbose, None).await +} + +pub async fn execute_with_runtime( + file_path: Option<&Path>, + args: Vec, + mounts: Vec, + json: bool, + verbose: bool, + shared_runtime: Option>, ) -> Result { 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(); diff --git a/crates/capsule-cli/src/commands/worker.rs b/crates/capsule-cli/src/commands/worker.rs new file mode 100644 index 00000000..cc03a226 --- /dev/null +++ b/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 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, + #[serde(default)] + pub mounts: Vec, +} + +#[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::(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, + mounts: Vec, + runtime: Arc, +) -> Result { + 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()), + } +} diff --git a/crates/capsule-cli/src/main.rs b/crates/capsule-cli/src/main.rs index 5c250784..943905e9 100644 --- a/crates/capsule-cli/src/main.rs +++ b/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?; diff --git a/crates/capsule-sdk/javascript/src/run.ts b/crates/capsule-sdk/javascript/src/run.ts index 9bf53c6b..030bdae1 100644 --- a/crates/capsule-sdk/javascript/src/run.ts +++ b/crates/capsule-sdk/javascript/src/run.ts @@ -5,11 +5,13 @@ * from third party applications. */ -import { execFile } from 'child_process'; -import { resolve, extname, join } from 'path'; +import { execFile, spawn, ChildProcess } from 'child_process'; +import { resolve, extname } from 'path'; import { existsSync, writeFileSync, unlinkSync } from 'fs'; import { tmpdir } from 'os'; +import { join } from 'path'; import { randomUUID } from 'crypto'; +import { createInterface } from 'readline'; import { HostRequest } from './task'; export interface RunnerOptions { @@ -35,11 +37,82 @@ export interface RunnerResult { } const WASM_EXTENSIONS = new Set(['.wasm']); -const ARGS_FILE_THRESHOLD = 8 * 1024; +const ARGS_FILE_THRESHOLD = 8 * 1024; // 8KB + + +interface PendingRequest { + resolve: (result: RunnerResult) => void; + reject: (err: Error) => void; +} + +let workerProcess: ChildProcess | null = null; +let workerCapsulePath: string | null = null; +const pending = new Map(); + +function getWorker(capsulePath: string): ChildProcess { + if (workerProcess && (workerCapsulePath !== capsulePath || workerProcess.exitCode !== null)) { + workerProcess.kill(); + workerProcess = null; + } + + if (!workerProcess) { + const command = getCapsuleCommand(capsulePath); + + let child: ChildProcess; + if (process.platform === 'win32') { + const comspec = process.env.comspec || 'cmd.exe'; + child = spawn(comspec, ['/d', '/s', '/c', command, 'worker'], { stdio: ['pipe', 'pipe', 'inherit'] }); + } else { + child = spawn(command, ['worker'], { stdio: ['pipe', 'pipe', 'inherit'] }); + } + + const rl = createInterface({ input: child.stdout! }); + rl.on('line', (line) => { + let response: { id: string; output?: unknown; error?: string }; + try { + response = JSON.parse(line); + } catch { + return; + } + + const request = pending.get(response.id); + if (!request) return; + pending.delete(response.id); + + if (response.error) { + request.reject(new Error(response.error)); + } else { + request.resolve(response.output as RunnerResult); + } + }); + + child.on('exit', () => { + for (const [id, req] of pending) { + req.reject(new Error('Capsule worker process exited unexpectedly')); + pending.delete(id); + } + workerProcess = null; + }); + + child.on('error', (err) => { + if ((err as NodeJS.ErrnoException).code === 'ENOENT') { + for (const [id, req] of pending) { + req.reject(new Error(`Capsule CLI not found. Use 'npm install -g @capsule-run/cli' to install it.`)); + pending.delete(id); + } + } + workerProcess = null; + }); + + workerProcess = child; + workerCapsulePath = capsulePath; + + process.once('exit', () => workerProcess?.kill()); + } + + return workerProcess; +} -/** - * Get the appropriate capsule command for the current platform - */ function getCapsuleCommand(capsulePath: string): string { if (process.platform === 'win32' && !capsulePath.endsWith('.cmd')) { return `${capsulePath}.cmd`; @@ -47,37 +120,43 @@ function getCapsuleCommand(capsulePath: string): string { return capsulePath; } -/** - * Write args to a temp file and return its path. - * Caller is responsible for deleting it. - */ function writeArgsFile(args: string[]): string { const path = join(tmpdir(), `capsule-args-${randomUUID()}.json`); writeFileSync(path, JSON.stringify(args), 'utf-8'); return path; } -/** - * Run a Capsule task from a third-party application - * - * @param options - Runner options - * @returns Promise with the runner result - */ -export function run(options: RunnerOptions): Promise { +// --- run() via persistent worker --- + +function runViaWorker(options: RunnerOptions): Promise { + const { file, args = [], mounts = [], capsulePath = 'capsule' } = options; + const id = randomUUID(); + console.time('runViaWorker' + id) + return new Promise((resolve, reject) => { + const worker = getWorker(capsulePath); + + pending.set(id, { resolve, reject }); + + const request = JSON.stringify({ id, file, args, mounts }) + '\n'; + worker.stdin!.write(request, (err) => { + if (err) { + pending.delete(id); + reject(new Error(`Failed to send task to worker: ${err.message}`)); + } + }); + console.timeEnd('runViaWorker' + id) + }); +} + +// --- run() via subprocess (fallback) --- + +function runViaSubprocess(options: RunnerOptions): Promise { const { file, args = [], mounts = [], cwd, capsulePath = 'capsule' } = options; const command = getCapsuleCommand(capsulePath); const resolvedFile = resolve(cwd || process.cwd(), file); const ext = extname(resolvedFile).toLowerCase(); const isWasm = WASM_EXTENSIONS.has(ext); - - if (!existsSync(resolvedFile)) { - const hint = isWasm - ? ` Run \`capsule build\` first to generate the .wasm artifact.` - : ''; - return Promise.reject(new Error(`File not found: ${resolvedFile}.${hint}`)); - } - const subcommand = isWasm ? 'exec' : 'run'; const mountFlags = mounts.flatMap(m => ['--mount', m]); @@ -112,9 +191,7 @@ export function run(options: RunnerOptions): Promise { if (error && !stdout) { if ((error as NodeJS.ErrnoException).code === 'ENOENT') { - reject(new Error( - `Capsule CLI not found. Use 'npm install -g @capsule-run/cli' to install it.` - )); + reject(new Error(`Capsule CLI not found. Use 'npm install -g @capsule-run/cli' to install it.`)); return; } reject(new Error(stderr || error.message)); @@ -123,11 +200,36 @@ export function run(options: RunnerOptions): Promise { try { const lines = stdout.trim().split('\n'); - const jsonLine = lines[lines.length - 1]; - resolve(JSON.parse(jsonLine)); + resolve(JSON.parse(lines[lines.length - 1])); } catch { reject(new Error(`Failed to parse Capsule output: ${stdout}`)); } }); }); } + +/** + * Run a Capsule task from a third-party application. + * Uses a persistent worker process to avoid per-call subprocess spawn overhead. + * + * @param options - Runner options + * @returns Promise with the runner result + */ +export async function run(options: RunnerOptions): Promise { + const { file, cwd } = options; + + const resolvedFile = resolve(cwd || process.cwd(), file); + const ext = extname(resolvedFile).toLowerCase(); + const isWasm = WASM_EXTENSIONS.has(ext); + + if (!existsSync(resolvedFile)) { + const hint = isWasm ? ` Run \`capsule build\` first to generate the .wasm artifact.` : ''; + throw new Error(`File not found: ${resolvedFile}.${hint}`); + } + + try { + return await runViaWorker(options); + } catch { + return runViaSubprocess(options); + } +} From cfb76349085ca3fcfde4138bf748baaac16e276b Mon Sep 17 00:00:00 2001 From: Mavdol Date: Fri, 17 Apr 2026 14:04:38 +0200 Subject: [PATCH 2/2] implement persistent worker process for low-latency task execution in Python SDK --- .../python/src/capsule/__init__.py | 1 + crates/capsule-sdk/python/src/capsule/run.py | 39 ++++ .../capsule-sdk/python/src/capsule/worker.py | 194 ++++++++++++++++++ 3 files changed, 234 insertions(+) create mode 100644 crates/capsule-sdk/python/src/capsule/worker.py diff --git a/crates/capsule-sdk/python/src/capsule/__init__.py b/crates/capsule-sdk/python/src/capsule/__init__.py index 4cf89ad8..98494a51 100644 --- a/crates/capsule-sdk/python/src/capsule/__init__.py +++ b/crates/capsule-sdk/python/src/capsule/__init__.py @@ -3,5 +3,6 @@ from . import http from .app import TaskRunner from .run import run, RunnerOptions, RunnerResult, ExecutionInfo, ErrorInfo +from .worker import run_with_worker, close_all exports = app.exports diff --git a/crates/capsule-sdk/python/src/capsule/run.py b/crates/capsule-sdk/python/src/capsule/run.py index ceaac533..fc53a6ab 100644 --- a/crates/capsule-sdk/python/src/capsule/run.py +++ b/crates/capsule-sdk/python/src/capsule/run.py @@ -63,15 +63,21 @@ async def run( ) -> RunnerResult: """Run a Capsule task from a third-party application. + Uses a persistent worker process for low latency (7-10ms per call). + Falls back to spawning a subprocess if the worker is unavailable. + Args: file: Path to the source file or pre-built .wasm artifact args: Arguments to pass to the task's main function + mounts: Mount specs (HOST[::GUEST][:ro|:rw]) cwd: Working directory (used to resolve relative paths) capsule_path: Path to the capsule CLI binary Returns: RunnerResult with task result and execution metadata """ + from .worker import run_with_worker, _mark_unavailable, _is_unavailable + args = args or [] mounts = mounts or [] @@ -93,6 +99,39 @@ async def run( ) raise FileNotFoundError(f"File not found: {resolved_file}") + if not _is_unavailable(capsule_path): + try: + return await run_with_worker( + file=file, + args=args, + mounts=mounts, + cwd=cwd, + capsule_path=capsule_path, + ) + except FileNotFoundError: + _mark_unavailable(capsule_path) + except Exception: + pass + + return await _run_subprocess( + resolved_file=resolved_file, + ext=ext, + args=args, + mounts=mounts, + cwd=cwd, + capsule_path=capsule_path, + ) + + +async def _run_subprocess( + *, + resolved_file: str, + ext: str, + args: list[str], + mounts: list[str], + cwd: Optional[str], + capsule_path: str, +) -> RunnerResult: subcommand = "exec" if ext in _WASM_EXTENSIONS else "run" mount_flags = [flag for m in mounts for flag in ("--mount", m)] diff --git a/crates/capsule-sdk/python/src/capsule/worker.py b/crates/capsule-sdk/python/src/capsule/worker.py new file mode 100644 index 00000000..8b1ad82c --- /dev/null +++ b/crates/capsule-sdk/python/src/capsule/worker.py @@ -0,0 +1,194 @@ +""" +Capsule SDK - Worker Client + +The worker process is created lazily on the first call and is shared for the +lifetime of the Python process. It is safe to call from multiple async tasks +concurrently — each call gets its own response future keyed by a unique id. +""" + +import asyncio +import json +import os +import uuid +from typing import Optional + +from .run import RunnerResult, _WASM_EXTENSIONS + + +class _WorkerClient: + """Manages a single long-lived `capsule worker` subprocess.""" + + def __init__(self, capsule_path: str, cwd: Optional[str]) -> None: + self._capsule_path = capsule_path + self._cwd = cwd + self._process: Optional[asyncio.subprocess.Process] = None + self._pending: dict[str, asyncio.Future[str]] = {} + self._reader_task: Optional[asyncio.Task] = None + self._lock = asyncio.Lock() + + async def _start(self) -> None: + self._process = await asyncio.create_subprocess_exec( + self._capsule_path, + "worker", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + cwd=self._cwd, + ) + self._reader_task = asyncio.create_task(self._reader_loop()) + + async def _reader_loop(self) -> None: + assert self._process and self._process.stdout + while True: + try: + line = await self._process.stdout.readline() + except Exception: + break + + if not line: + break + try: + response = json.loads(line.decode("utf-8")) + req_id = response.get("id") + future = self._pending.pop(req_id, None) + if future and not future.done(): + future.set_result(line.decode("utf-8")) + except Exception: + continue + + for future in self._pending.values(): + if not future.done(): + future.set_exception(RuntimeError("capsule worker process exited unexpectedly")) + + self._pending.clear() + + async def _ensure_running(self) -> None: + async with self._lock: + dead = ( + self._process is None + or self._process.returncode is not None + ) + if dead: + await self._start() + + async def send(self, file: str, args: list[str], mounts: list[str]) -> str: + await self._ensure_running() + + req_id = uuid.uuid4().hex + request = json.dumps({"id": req_id, "file": file, "args": args, "mounts": mounts}) + + loop = asyncio.get_event_loop() + future: asyncio.Future[str] = loop.create_future() + self._pending[req_id] = future + + assert self._process and self._process.stdin + self._process.stdin.write((request + "\n").encode("utf-8")) + await self._process.stdin.drain() + + return await future + + async def close(self) -> None: + if self._process and self._process.returncode is None: + try: + assert self._process.stdin + self._process.stdin.close() + await self._process.wait() + except Exception: + pass + if self._reader_task: + self._reader_task.cancel() + + +_clients: dict[tuple[str, Optional[str]], _WorkerClient] = {} +_clients_lock: Optional[asyncio.Lock] = None + +# Tracks capsule_path values where the binary was confirmed missing +_unavailable: set[str] = set() + + +def _mark_unavailable(capsule_path: str) -> None: + _unavailable.add(capsule_path) + + +def _is_unavailable(capsule_path: str) -> bool: + return capsule_path in _unavailable + + +def _get_lock() -> asyncio.Lock: + global _clients_lock + if _clients_lock is None: + _clients_lock = asyncio.Lock() + return _clients_lock + + +async def _get_client(capsule_path: str, cwd: Optional[str]) -> _WorkerClient: + key = (capsule_path, cwd) + async with _get_lock(): + if key not in _clients: + _clients[key] = _WorkerClient(capsule_path, cwd) + return _clients[key] + + +async def run_with_worker( + *, + file: str, + args: Optional[list[str]] = None, + mounts: Optional[list[str]] = None, + cwd: Optional[str] = None, + capsule_path: str = "capsule", +) -> RunnerResult: + """Run a Capsule task via a persistent worker process (7-10ms per call). + + Uses a single long-lived `capsule worker` subprocess per combination. + + Args: + file: Path to the source file or pre-built .wasm artifact + args: Arguments to pass to the task's main function + mounts: Mount specs forwarded to the worker (HOST[::GUEST][:ro|:rw]) + cwd: Working directory for resolving relative paths + capsule_path: Path to the capsule CLI binary + + Returns: + RunnerResult with task result and execution metadata + """ + args = args or [] + mounts = mounts or [] + + resolved_file = os.path.abspath(os.path.join(cwd or os.getcwd(), file)) + ext = os.path.splitext(resolved_file)[1].lower() + + if not os.path.exists(resolved_file): + if ext in _WASM_EXTENSIONS: + raise FileNotFoundError( + f"Artifact not found: {resolved_file}. " + "Run `capsule build` first to generate the .wasm artifact." + ) + raise FileNotFoundError(f"File not found: {resolved_file}") + + client = await _get_client(capsule_path, cwd) + + try: + raw = await client.send(resolved_file, args, mounts) + except RuntimeError: + key = (capsule_path, cwd) + _clients.pop(key, None) + raise + + try: + lines = raw.strip().split("\n") + response = json.loads(lines[-1]) + except (json.JSONDecodeError, IndexError): + raise ValueError(f"Failed to parse worker response: {raw}") + + if "error" in response and "output" not in response: + raise RuntimeError(response["error"]) + + return response.get("output", response) + + +async def close_all() -> None: + """Close all active worker processes. Call on application shutdown.""" + for client in list(_clients.values()): + await client.close() + + _clients.clear()