From 962f7e609454cd2d54e3ba00077ede7c3bcab7ff Mon Sep 17 00:00:00 2001 From: Navneet Chaudhary Date: Mon, 17 Aug 2026 16:10:21 +0530 Subject: [PATCH 1/3] feat: add async command jobs --- Cargo.lock | 2 + Cargo.toml | 6 + README.md | 25 +- src/command.rs | 84 +- src/command_jobs.rs | 1268 +++++++++++++++++++++++++++++ src/main.rs | 16 +- src/mcp.rs | 954 +++++++++++++++++++++- src/process_runner.rs | 594 ++++++++++++++ src/server.rs | 122 +++ src/state.rs | 3 + src/widget/catdesk_dashboard.html | 16 +- 11 files changed, 3004 insertions(+), 86 deletions(-) create mode 100644 src/command_jobs.rs create mode 100644 src/process_runner.rs diff --git a/Cargo.lock b/Cargo.lock index 0a0b8b1..8e439aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -392,6 +392,7 @@ dependencies = [ "globset", "ignore", "image", + "libc", "ngrok", "rand 0.8.5", "rand_mt", @@ -408,6 +409,7 @@ dependencies = [ "tree-sitter", "tree-sitter-bash", "uuid", + "windows-sys 0.61.2", "zip", ] diff --git a/Cargo.toml b/Cargo.toml index b5e59c9..955539a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,12 @@ regex = "1" globset = "0.4" ignore = "0.4" +[target.'cfg(unix)'.dependencies] +libc = "0.2" + +[target.'cfg(windows)'.dependencies] +windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Security", "Win32_System_JobObjects", "Win32_System_Threading"] } + [profile.release] lto = true strip = true diff --git a/README.md b/README.md index 6964122..8043654 100644 --- a/README.md +++ b/README.md @@ -160,19 +160,24 @@ CatDesk is a coding tool and a custom connector. Always use CatDesk if the user # Tools -CatDesk has two local tool modes: `multi-tools` exposes 7 tools, and `read-only` exposes 3 tools. +CatDesk has two local tool modes: `multi-tools` exposes 10 tools, and `read-only` exposes 3 tools. CatDesk's local tools in `multi-tools` mode are: -| Tool | Type | What it does | -| --------------------- | ----- | ------------------------------------------------------------- | -| `catdesk_instruction` | Guide | Returns CatDesk usage instructions and render Binagotchy | -| `read` | Read | Reads a text file from the workspace | -| `search` | Read | Searches workspace text with `rg`, `grep`, or built-in search | -| `write` | Write | Creates or overwrites a file | -| `edit` | Write | Replaces exact text inside a file | -| `delete` | Write | Deletes a file or directory | -| `run_command` | Shell | Runs a shell command inside the workspace. | +| Tool | Type | What it does | +| --------------------- | ----- | -------------------------------------------------------------------------- | +| `catdesk_instruction` | Guide | Returns CatDesk usage instructions and render Binagotchy | +| `read` | Read | Reads a text file from the workspace | +| `search` | Read | Searches workspace text with `rg`, `grep`, or built-in search | +| `write` | Write | Creates or overwrites a file | +| `edit` | Write | Replaces exact text inside a file | +| `delete` | Write | Deletes a file or directory | +| `run_command` | Shell | Runs a short shell command and waits for completion | +| `start_command` | Job | Starts a long-running shell command and immediately returns a job ID | +| `poll_command` | Job | Reads incremental output and status from a background command | +| `cancel_command` | Job | Stops a background command and its child process tree | + +Long-running commands are deliberately decoupled from the lifetime of an MCP HTTP request. Builds, compilation, dependency installation, long test suites, and development servers should use `start_command`, then `poll_command` with the returned cursor. Poll responses are bounded; if `hasMoreOutput` is true, keep polling with `nextCursor` even after the command reaches a terminal state to drain the remaining buffered output. `run_command` remains the simpler path for short commands and has a 120-second maximum timeout. If browser mode is enabled, CatDesk can also expose extra browser/devtools tools. Those are provided by the browser bridge, so the exact list depends on your environment. diff --git a/src/command.rs b/src/command.rs index 3ac2a16..cc7ca66 100644 --- a/src/command.rs +++ b/src/command.rs @@ -1,13 +1,10 @@ use std::path::{Path, PathBuf}; -use std::time::Instant; -use tokio::process::Command; -use tokio::time::{Duration, timeout}; use tree_sitter::{Node, Parser}; use tree_sitter_bash::LANGUAGE as BASH_LANGUAGE; const MAX_BUFFER_BYTES: usize = 1024 * 1024; -const DEFAULT_TIMEOUT_MS: u64 = 30_000; -const MAX_TIMEOUT_MS: u64 = 120_000; +pub const DEFAULT_TIMEOUT_MS: u64 = 30_000; +pub const MAX_TIMEOUT_MS: u64 = 120_000; pub const CATDESK_CO_AUTHOR_TRAILER: &str = "Co-Authored-By: CatDesk"; #[derive(Debug)] @@ -15,7 +12,11 @@ pub struct CommandResult { pub stdout: String, pub stderr: String, pub success: bool, + pub exit_code: Option, pub elapsed_ms: u64, + pub timed_out: bool, + pub stdout_truncated: bool, + pub stderr_truncated: bool, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -170,66 +171,27 @@ pub fn detect_move_path_intercept(command: &str) -> Option CommandResult { - let start = Instant::now(); - let mut shell = shell_command(command); - let fut = shell.current_dir(cwd).output(); - - match timeout(Duration::from_millis(timeout_ms), fut).await { - Ok(Ok(output)) => { - let elapsed_ms = start.elapsed().as_millis() as u64; - let stdout = String::from_utf8_lossy( - &output.stdout[..output.stdout.len().min(MAX_BUFFER_BYTES)], - ) - .to_string(); - let stderr = String::from_utf8_lossy( - &output.stderr[..output.stderr.len().min(MAX_BUFFER_BYTES)], - ) - .to_string(); - CommandResult { - stdout, - stderr, - success: output.status.success(), - elapsed_ms, - } - } - Ok(Err(e)) => CommandResult { - stdout: String::new(), - stderr: format!("Failed to execute: {e}"), - success: false, - elapsed_ms: start.elapsed().as_millis() as u64, - }, - Err(_) => CommandResult { - stdout: String::new(), - stderr: format!("Command timed out after {timeout_ms} ms"), - success: false, - elapsed_ms: start.elapsed().as_millis() as u64, - }, + let result = + crate::process_runner::run_shell_command(command, cwd, timeout_ms, MAX_BUFFER_BYTES).await; + + CommandResult { + stdout: result.stdout, + stderr: result.stderr, + success: result.success, + exit_code: result.exit_code, + elapsed_ms: result.elapsed_ms, + timed_out: result.timed_out, + stdout_truncated: result.stdout_truncated, + stderr_truncated: result.stderr_truncated, } } -#[cfg(windows)] -fn shell_command(command: &str) -> Command { - let mut shell = Command::new("powershell.exe"); - shell - .arg("-NoLogo") - .arg("-NoProfile") - .arg("-NonInteractive") - .arg("-ExecutionPolicy") - .arg("Bypass") - .arg("-Command") - .arg(command); - shell -} - -#[cfg(not(windows))] -fn shell_command(command: &str) -> Command { - let mut shell = Command::new("/bin/bash"); - shell.arg("-c").arg(command); - shell -} - /// Format stdout+stderr into a single string. pub fn format_result(r: &CommandResult) -> String { let mut out = String::new(); diff --git a/src/command_jobs.rs b/src/command_jobs.rs new file mode 100644 index 0000000..232200b --- /dev/null +++ b/src/command_jobs.rs @@ -0,0 +1,1268 @@ +use std::collections::{HashMap, HashSet, VecDeque}; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration as StdDuration, Instant}; + +use serde::Serialize; +use tokio::io::{AsyncRead, AsyncReadExt}; +use tokio::sync::{Mutex, Notify, RwLock, watch}; +use tokio::time::{Duration, timeout}; +use uuid::Uuid; + +use crate::process_runner; + +pub const DEFAULT_JOB_TIMEOUT_MS: u64 = 30 * 60 * 1_000; +pub const MAX_JOB_TIMEOUT_MS: u64 = 24 * 60 * 60 * 1_000; +pub const MAX_POLL_WAIT_MS: u64 = 30_000; +const MAX_ACTIVE_JOBS: usize = 8; +const MAX_RETAINED_JOBS: usize = 64; +const TERMINAL_JOB_TTL: StdDuration = StdDuration::from_secs(60 * 60); +const IDEMPOTENCY_WINDOW: StdDuration = StdDuration::from_secs(30); +const MAX_OUTPUT_BYTES_PER_JOB: usize = 4 * 1024 * 1024; +const MAX_POLL_OUTPUT_BYTES: usize = 128 * 1024; +const READ_CHUNK_BYTES: usize = 8 * 1024; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum CommandJobState { + Running, + Succeeded, + Failed, + Cancelled, + TimedOut, +} + +impl CommandJobState { + pub fn as_str(self) -> &'static str { + match self { + Self::Running => "running", + Self::Succeeded => "succeeded", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + Self::TimedOut => "timed_out", + } + } + + pub fn is_terminal(self) -> bool { + !matches!(self, Self::Running) + } +} + +#[derive(Clone, Debug, Serialize)] +pub struct CommandOutputEvent { + pub seq: u64, + pub stream: &'static str, + pub text: String, +} + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CommandJobSnapshot { + pub job_id: String, + pub command: String, + pub cwd: String, + pub state: CommandJobState, + pub elapsed_ms: u64, + pub exit_code: Option, + pub events: Vec, + pub next_cursor: u64, + pub has_more_output: bool, + pub output_truncated: bool, + pub timeout_ms: u64, +} + +#[derive(Clone, Debug)] +pub struct StartCommandResult { + pub snapshot: CommandJobSnapshot, + pub deduplicated: bool, +} + +#[derive(Debug)] +struct JobRuntime { + state: CommandJobState, + exit_code: Option, + finished_at: Option, + events: VecDeque, + retained_output_bytes: usize, + next_seq: u64, + output_truncated: bool, +} + +impl Default for JobRuntime { + fn default() -> Self { + Self { + state: CommandJobState::Running, + exit_code: None, + finished_at: None, + events: VecDeque::new(), + retained_output_bytes: 0, + next_seq: 1, + output_truncated: false, + } + } +} + +#[derive(Debug)] +struct CommandJob { + id: String, + command: String, + cwd: PathBuf, + started_at: Instant, + timeout_ms: u64, + runtime: Mutex, + changed: Notify, + cancel_tx: watch::Sender, +} + +impl CommandJob { + fn new(command: String, cwd: PathBuf, timeout_ms: u64) -> (Arc, watch::Receiver) { + let (cancel_tx, cancel_rx) = watch::channel(false); + ( + Arc::new(Self { + id: Uuid::new_v4().to_string(), + command, + cwd, + started_at: Instant::now(), + timeout_ms, + runtime: Mutex::new(JobRuntime::default()), + changed: Notify::new(), + cancel_tx, + }), + cancel_rx, + ) + } + + async fn append_output(&self, stream: &'static str, bytes: &[u8]) { + if bytes.is_empty() { + return; + } + let text = String::from_utf8_lossy(bytes).into_owned(); + let event_bytes = text.len(); + let mut runtime = self.runtime.lock().await; + let seq = runtime.next_seq; + runtime.next_seq = runtime.next_seq.saturating_add(1); + runtime + .events + .push_back(CommandOutputEvent { seq, stream, text }); + runtime.retained_output_bytes = runtime.retained_output_bytes.saturating_add(event_bytes); + + while runtime.retained_output_bytes > MAX_OUTPUT_BYTES_PER_JOB { + let Some(removed) = runtime.events.pop_front() else { + break; + }; + runtime.retained_output_bytes = runtime + .retained_output_bytes + .saturating_sub(removed.text.len()); + runtime.output_truncated = true; + } + drop(runtime); + self.changed.notify_waiters(); + } + + async fn finish(&self, state: CommandJobState, exit_code: Option) { + let mut runtime = self.runtime.lock().await; + if runtime.state.is_terminal() { + return; + } + runtime.state = state; + runtime.exit_code = exit_code; + runtime.finished_at = Some(Instant::now()); + drop(runtime); + self.changed.notify_waiters(); + } + + async fn snapshot(&self, after: u64) -> CommandJobSnapshot { + let runtime = self.runtime.lock().await; + let first_retained_seq = runtime + .events + .front() + .map(|event| event.seq) + .unwrap_or(runtime.next_seq); + let cursor_fell_behind = after.saturating_add(1) < first_retained_seq; + let latest_cursor = runtime.next_seq.saturating_sub(1); + let mut events = Vec::new(); + let mut response_bytes = 0usize; + for event in runtime.events.iter().filter(|event| event.seq > after) { + let event_bytes = event.text.len(); + if !events.is_empty() + && response_bytes.saturating_add(event_bytes) > MAX_POLL_OUTPUT_BYTES + { + break; + } + response_bytes = response_bytes.saturating_add(event_bytes); + events.push(event.clone()); + } + let next_cursor = events + .last() + .map(|event| event.seq) + .unwrap_or(latest_cursor); + let has_more_output = next_cursor < latest_cursor; + CommandJobSnapshot { + job_id: self.id.clone(), + command: self.command.clone(), + cwd: self.cwd.to_string_lossy().into_owned(), + state: runtime.state, + elapsed_ms: self.started_at.elapsed().as_millis() as u64, + exit_code: runtime.exit_code, + events, + next_cursor, + has_more_output, + output_truncated: runtime.output_truncated || cursor_fell_behind, + timeout_ms: self.timeout_ms, + } + } + + async fn terminal_age(&self) -> Option { + self.runtime + .lock() + .await + .finished_at + .map(|finished| finished.elapsed()) + } +} + +#[derive(Default)] +struct ManagerState { + jobs: HashMap>, + // Retry dedupe is intentionally short-lived. JSON-RPC request IDs are only + // correlation IDs and may be reused later by a stateless client. + request_jobs: HashMap, +} + +#[derive(Clone, Default)] +pub struct CommandJobManager { + inner: Arc>, + // Starting a job performs a dedupe lookup, active-job capacity check, and + // registry insertion. Serialize that short critical section so concurrent + // MCP requests cannot both pass the checks and create duplicate/overflow jobs. + start_lock: Arc>, + // App shutdown is terminal for this manager. Once set, no new background + // command may be created even if an MCP request races with shutdown. + shutting_down: Arc, +} + +impl CommandJobManager { + pub fn new() -> Self { + Self::default() + } + + pub fn normalize_timeout(timeout_ms: Option) -> Result { + match timeout_ms { + None => Ok(DEFAULT_JOB_TIMEOUT_MS), + Some(0) => Err("timeout must be at least 1 ms".to_string()), + Some(value) if value > MAX_JOB_TIMEOUT_MS => Err(format!( + "timeout exceeds the maximum background command runtime of {MAX_JOB_TIMEOUT_MS} ms" + )), + Some(value) => Ok(value), + } + } + + pub async fn start( + &self, + command: String, + cwd: PathBuf, + timeout_ms: u64, + request_key: Option, + ) -> Result { + let _start_guard = self.start_lock.lock().await; + if self.shutting_down.load(Ordering::Acquire) { + return Err( + "command job manager is shutting down; new commands are not accepted".to_string(), + ); + } + self.cleanup().await; + + if let Some(key) = request_key.as_deref() { + let existing = { + let manager = self.inner.read().await; + manager + .request_jobs + .get(key) + .and_then(|(job_id, created_at)| { + if created_at.elapsed() <= IDEMPOTENCY_WINDOW { + manager.jobs.get(job_id).cloned() + } else { + None + } + }) + }; + if let Some(job) = existing { + if job.command != command || job.cwd != cwd || job.timeout_ms != timeout_ms { + return Err( + "the same MCP request id was reused with different start_command arguments" + .to_string(), + ); + } + return Ok(StartCommandResult { + snapshot: job.snapshot(0).await, + deduplicated: true, + }); + } + } + + let active_count = { + let jobs = { + let manager = self.inner.read().await; + manager.jobs.values().cloned().collect::>() + }; + let mut active = 0usize; + for job in jobs { + if job.runtime.lock().await.state == CommandJobState::Running { + active += 1; + } + } + active + }; + if active_count >= MAX_ACTIVE_JOBS { + return Err(format!( + "too many active command jobs ({active_count}); maximum is {MAX_ACTIVE_JOBS}. Poll or cancel an existing job before starting another" + )); + } + + let (job, cancel_rx) = CommandJob::new(command, cwd, timeout_ms); + let job_id = job.id.clone(); + { + let mut manager = self.inner.write().await; + manager.jobs.insert(job_id.clone(), job.clone()); + if let Some(key) = request_key { + manager + .request_jobs + .insert(key, (job_id.clone(), Instant::now())); + } + } + + tokio::spawn(run_job(job.clone(), cancel_rx)); + Ok(StartCommandResult { + snapshot: job.snapshot(0).await, + deduplicated: false, + }) + } + + pub async fn poll( + &self, + job_id: &str, + after: u64, + wait_ms: u64, + ) -> Result { + self.cleanup().await; + let job = self.get_job(job_id).await?; + let wait_ms = wait_ms.min(MAX_POLL_WAIT_MS); + + // Register the notification future before the first snapshot so output + // arriving between the check and wait cannot be missed. + let notified = job.changed.notified(); + let snapshot = job.snapshot(after).await; + if snapshot.state.is_terminal() || !snapshot.events.is_empty() || wait_ms == 0 { + return Ok(snapshot); + } + + let _ = timeout(Duration::from_millis(wait_ms), notified).await; + Ok(job.snapshot(after).await) + } + + pub async fn cancel(&self, job_id: &str) -> Result { + self.cleanup().await; + let job = self.get_job(job_id).await?; + + let current = job.snapshot(0).await; + if current.state.is_terminal() { + return Ok(current); + } + let _ = job.cancel_tx.send(true); + + // Cancellation itself remains a short MCP operation, but ordinary + // stdout/stderr notifications must not make cancel_command return a + // misleading Running state. Wait until terminal or the bounded deadline. + let deadline = Instant::now() + StdDuration::from_secs(5); + loop { + // Register before the snapshot so a terminal transition cannot land + // in the check/wait gap. + let notified = job.changed.notified(); + let snapshot = job.snapshot(0).await; + if snapshot.state.is_terminal() { + return Ok(snapshot); + } + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Ok(snapshot); + } + if timeout(remaining, notified).await.is_err() { + return Ok(job.snapshot(0).await); + } + } + } + + async fn get_job(&self, job_id: &str) -> Result, String> { + self.inner + .read() + .await + .jobs + .get(job_id) + .cloned() + .ok_or_else(|| format!("unknown or expired command job: {job_id}")) + } + + /// Cancel every command still owned by CatDesk and wait briefly for the + /// runners to terminate their process trees. Used during application exit; + /// ordinary MCP request completion deliberately does not call this. + pub async fn cancel_all(&self) { + // Serialize with start(): either a start completes before this guard and + // is included below, or shutdown wins and that start is rejected. + let _start_guard = self.start_lock.lock().await; + self.shutting_down.store(true, Ordering::Release); + + let jobs = { + let manager = self.inner.read().await; + manager.jobs.values().cloned().collect::>() + }; + for job in &jobs { + if job.runtime.lock().await.state == CommandJobState::Running { + let _ = job.cancel_tx.send(true); + } + } + + let deadline = Instant::now() + StdDuration::from_secs(5); + loop { + let mut any_running = false; + for job in &jobs { + if job.runtime.lock().await.state == CommandJobState::Running { + any_running = true; + break; + } + } + if !any_running || Instant::now() >= deadline { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + pub async fn cleanup(&self) { + let jobs = { + let manager = self.inner.read().await; + manager + .jobs + .iter() + .map(|(id, job)| (id.clone(), job.clone())) + .collect::>() + }; + + let mut expired = Vec::new(); + let mut terminal = Vec::new(); + for (id, job) in jobs { + if let Some(age) = job.terminal_age().await { + if age >= TERMINAL_JOB_TTL { + expired.push(id); + } else { + terminal.push((id, age)); + } + } + } + + terminal.sort_by_key(|(_, age)| std::cmp::Reverse(*age)); + let retained_count = { + let manager = self.inner.read().await; + manager.jobs.len().saturating_sub(expired.len()) + }; + if retained_count > MAX_RETAINED_JOBS { + let overflow = retained_count - MAX_RETAINED_JOBS; + expired.extend(terminal.into_iter().take(overflow).map(|(id, _)| id)); + } + let mut manager = self.inner.write().await; + for id in &expired { + manager.jobs.remove(id); + } + let live_job_ids = manager.jobs.keys().cloned().collect::>(); + manager.request_jobs.retain(|_, (job_id, created_at)| { + created_at.elapsed() <= IDEMPOTENCY_WINDOW && live_job_ids.contains(job_id) + }); + } +} + +fn decode_utf8_incremental( + pending: &mut Vec, + chunk: &[u8], + end_of_stream: bool, +) -> Vec { + pending.extend_from_slice(chunk); + let mut decoded = Vec::new(); + + loop { + match std::str::from_utf8(pending) { + Ok(text) => { + if !text.is_empty() { + decoded.push(text.to_string()); + } + pending.clear(); + break; + } + Err(error) => { + let valid_up_to = error.valid_up_to(); + if valid_up_to > 0 { + let text = String::from_utf8_lossy(&pending[..valid_up_to]).into_owned(); + decoded.push(text); + pending.drain(..valid_up_to); + continue; + } + match error.error_len() { + Some(invalid_len) => { + decoded.push("๏ฟฝ".to_string()); + pending.drain(..invalid_len.min(pending.len())); + } + None => break, // incomplete UTF-8 sequence; keep it for the next read + } + } + } + } + + if end_of_stream && !pending.is_empty() { + decoded.push(String::from_utf8_lossy(pending).into_owned()); + pending.clear(); + } + decoded +} + +async fn read_job_output(job: Arc, stream: &'static str, mut reader: R) +where + R: AsyncRead + Unpin, +{ + let mut buffer = vec![0_u8; READ_CHUNK_BYTES]; + let mut pending_utf8 = Vec::with_capacity(4); + loop { + match reader.read(&mut buffer).await { + Ok(0) => { + for text in decode_utf8_incremental(&mut pending_utf8, &[], true) { + job.append_output(stream, text.as_bytes()).await; + } + break; + } + Ok(read) => { + for text in decode_utf8_incremental(&mut pending_utf8, &buffer[..read], false) { + job.append_output(stream, text.as_bytes()).await; + } + } + Err(error) => { + for text in decode_utf8_incremental(&mut pending_utf8, &[], true) { + job.append_output(stream, text.as_bytes()).await; + } + job.append_output( + "stderr", + format!("CatDesk failed to read {stream}: {error}\n").as_bytes(), + ) + .await; + break; + } + } + } +} + +async fn run_job(job: Arc, mut cancel_rx: watch::Receiver) { + let mut process = match process_runner::spawn_shell_command(&job.command, &job.cwd) { + Ok(process) => process, + Err(error) => { + job.append_output("stderr", format!("Failed to execute: {error}\n").as_bytes()) + .await; + job.finish(CommandJobState::Failed, None).await; + return; + } + }; + + let stdout_task = process + .take_stdout() + .map(|stdout| tokio::spawn(read_job_output(job.clone(), "stdout", stdout))); + let stderr_task = process + .take_stderr() + .map(|stderr| tokio::spawn(read_job_output(job.clone(), "stderr", stderr))); + + enum Completion { + Exited(std::io::Result), + Cancelled, + TimedOut, + } + + let completion = tokio::select! { + status = process.wait() => Completion::Exited(status), + _ = cancel_rx.changed() => Completion::Cancelled, + _ = tokio::time::sleep(Duration::from_millis(job.timeout_ms)) => Completion::TimedOut, + }; + + let (state, exit_code) = match completion { + Completion::Exited(Ok(status)) => { + process.disarm(); + if status.success() { + (CommandJobState::Succeeded, status.code()) + } else { + (CommandJobState::Failed, status.code()) + } + } + Completion::Exited(Err(error)) => { + process.terminate_tree(); + let _ = process.wait().await; + job.append_output( + "stderr", + format!("CatDesk failed while waiting for command: {error}\n").as_bytes(), + ) + .await; + (CommandJobState::Failed, None) + } + Completion::Cancelled => { + process.terminate_tree(); + let status = process.wait().await.ok(); + ( + CommandJobState::Cancelled, + status.and_then(|value| value.code()), + ) + } + Completion::TimedOut => { + process.terminate_tree(); + let status = process.wait().await.ok(); + job.append_output( + "stderr", + format!("Command timed out after {} ms\n", job.timeout_ms).as_bytes(), + ) + .await; + ( + CommandJobState::TimedOut, + status.and_then(|value| value.code()), + ) + } + }; + + if let Some(task) = stdout_task { + let _ = task.await; + } + if let Some(task) = stderr_task { + let _ = task.await; + } + job.finish(state, exit_code).await; +} + +#[cfg(test)] +mod tests { + use super::*; + + fn workspace(name: &str) -> PathBuf { + let path = std::env::temp_dir().join(format!("catdesk-jobs-{name}-{}", Uuid::new_v4())); + std::fs::create_dir_all(&path).expect("create test workspace"); + path + } + + async fn wait_terminal(manager: &CommandJobManager, job_id: &str) -> CommandJobSnapshot { + let mut cursor = 0; + for _ in 0..30 { + let snapshot = manager.poll(job_id, cursor, 250).await.expect("poll job"); + cursor = snapshot.next_cursor; + if snapshot.state.is_terminal() { + // Fetch from zero once terminal so callers that assert on output + // see the complete retained log rather than only the final delta. + return manager.poll(job_id, 0, 0).await.expect("read terminal job"); + } + } + panic!("job did not reach terminal state"); + } + + async fn wait_for_file(path: &std::path::Path) { + let deadline = Instant::now() + StdDuration::from_secs(5); + while !path.exists() && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!( + path.exists(), + "command never reached ready state: {}", + path.display() + ); + } + + #[tokio::test] + async fn background_job_returns_immediately_and_completes() { + let root = workspace("complete"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 300; Write-Output done" + } else { + "sleep 0.3; printf 'done\\n'" + }; + let started = Instant::now(); + let started_job = manager + .start(command.to_string(), root.clone(), 5_000, None) + .await + .expect("start job"); + assert!(started.elapsed() < StdDuration::from_millis(250)); + assert_eq!(started_job.snapshot.state, CommandJobState::Running); + + let snapshot = wait_terminal(&manager, &started_job.snapshot.job_id).await; + assert_eq!(snapshot.state, CommandJobState::Succeeded); + let text = snapshot + .events + .iter() + .map(|event| event.text.as_str()) + .collect::(); + assert!(text.contains("done")); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn polling_is_incremental_by_cursor() { + let root = workspace("cursor"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Write-Output first; Start-Sleep -Milliseconds 250; Write-Output second" + } else { + "printf 'first\\n'; sleep 0.25; printf 'second\\n'" + }; + let started = manager + .start(command.to_string(), root.clone(), 5_000, None) + .await + .expect("start job"); + + let deadline = Instant::now() + StdDuration::from_secs(5); + let first = loop { + let snapshot = manager + .poll(&started.snapshot.job_id, 0, 250) + .await + .expect("first poll"); + if !snapshot.events.is_empty() { + break snapshot; + } + assert!( + !snapshot.state.is_terminal(), + "job completed without producing expected first output" + ); + assert!( + Instant::now() < deadline, + "timed out waiting for first output" + ); + }; + + let first_cursor = first.next_cursor; + let deadline = Instant::now() + StdDuration::from_secs(5); + let second = loop { + let snapshot = manager + .poll(&started.snapshot.job_id, first_cursor, 250) + .await + .expect("second poll"); + if !snapshot.events.is_empty() || snapshot.state.is_terminal() { + break snapshot; + } + assert!( + Instant::now() < deadline, + "timed out waiting for second output" + ); + }; + assert!( + second.events.iter().all(|event| event.seq > first_cursor), + "incremental poll repeated an already-consumed event" + ); + let _ = wait_terminal(&manager, &started.snapshot.job_id).await; + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn cancellation_prevents_later_side_effects() { + let root = workspace("cancel"); + let ready = root.join("ready.txt"); + let sentinel = root.join("sentinel.txt"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Set-Content ready.txt ready; Start-Sleep -Seconds 3; Set-Content sentinel.txt survived" + } else { + "printf ready > ready.txt; sleep 3; printf survived > sentinel.txt" + }; + let started = manager + .start(command.to_string(), root.clone(), 10_000, None) + .await + .expect("start job"); + wait_for_file(&ready).await; + let cancelled = manager + .cancel(&started.snapshot.job_id) + .await + .expect("cancel job"); + assert!(matches!( + cancelled.state, + CommandJobState::Cancelled | CommandJobState::Running + )); + let terminal = wait_terminal(&manager, &started.snapshot.job_id).await; + assert_eq!(terminal.state, CommandJobState::Cancelled); + tokio::time::sleep(Duration::from_millis(900)).await; + assert!( + !sentinel.exists(), + "cancelled process survived and wrote sentinel" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn cancel_waits_for_terminal_state_despite_output_notifications() { + let root = workspace("cancel-terminal"); + let ready = root.join("ready.txt"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Set-Content ready.txt ready; 1..200 | ForEach-Object { Write-Output $_; Start-Sleep -Milliseconds 5 }; Start-Sleep -Seconds 3" + } else { + "printf ready > ready.txt; for i in $(seq 1 200); do printf '%s\\n' \"$i\"; sleep 0.005; done; sleep 3" + }; + let started = manager + .start(command.to_string(), root.clone(), 10_000, None) + .await + .expect("start noisy job"); + wait_for_file(&ready).await; + let cancelled = manager + .cancel(&started.snapshot.job_id) + .await + .expect("cancel noisy job"); + assert_eq!( + cancelled.state, + CommandJobState::Cancelled, + "cancel_command should wait past output notifications for terminal acknowledgement" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn cancellation_terminates_descendant_process_tree() { + let root = workspace("descendant-cancel"); + let sentinel = root.join("descendant.txt"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Start-Process powershell.exe -ArgumentList '-NoProfile','-Command','Start-Sleep -Milliseconds 800; Set-Content -Path descendant.txt -Value survived' -WorkingDirectory .; Start-Sleep -Seconds 5" + } else { + "(sleep 0.8; printf survived > descendant.txt) & sleep 5" + }; + let started = manager + .start(command.to_string(), root.clone(), 5_000, None) + .await + .expect("start descendant job"); + tokio::time::sleep(Duration::from_millis(150)).await; + let _ = manager + .cancel(&started.snapshot.job_id) + .await + .expect("cancel descendant job"); + let terminal = wait_terminal(&manager, &started.snapshot.job_id).await; + assert_eq!(terminal.state, CommandJobState::Cancelled); + tokio::time::sleep(Duration::from_millis(1_000)).await; + assert!( + !sentinel.exists(), + "cancelled root shell left a descendant process alive" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn job_timeout_terminates_process_tree() { + let root = workspace("timeout"); + let sentinel = root.join("sentinel.txt"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 800; Set-Content sentinel.txt survived" + } else { + "sleep 0.8; printf survived > sentinel.txt" + }; + let started = manager + .start(command.to_string(), root.clone(), 100, None) + .await + .expect("start job"); + let terminal = wait_terminal(&manager, &started.snapshot.job_id).await; + assert_eq!(terminal.state, CommandJobState::TimedOut); + tokio::time::sleep(Duration::from_millis(900)).await; + assert!( + !sentinel.exists(), + "timed-out job survived and wrote sentinel" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn incremental_utf8_decoder_preserves_split_multibyte_characters() { + let bytes = "build โœ“ ๐Ÿš€".as_bytes(); + let split = bytes.len() - 2; + let mut pending = Vec::new(); + let first = decode_utf8_incremental(&mut pending, &bytes[..split], false); + let second = decode_utf8_incremental(&mut pending, &bytes[split..], false); + let final_chunk = decode_utf8_incremental(&mut pending, &[], true); + let decoded = first + .into_iter() + .chain(second) + .chain(final_chunk) + .collect::(); + assert_eq!(decoded, "build โœ“ ๐Ÿš€"); + assert!(pending.is_empty()); + } + + #[tokio::test] + async fn cancel_all_terminates_active_jobs() { + let root = workspace("cancel-all"); + let ready = root.join("ready.txt"); + let sentinel = root.join("sentinel.txt"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Set-Content ready.txt ready; Start-Sleep -Seconds 3; Set-Content sentinel.txt survived" + } else { + "printf ready > ready.txt; sleep 3; printf survived > sentinel.txt" + }; + let started = manager + .start(command.to_string(), root.clone(), 10_000, None) + .await + .expect("start job"); + wait_for_file(&ready).await; + manager.cancel_all().await; + let terminal = manager + .poll(&started.snapshot.job_id, 0, 0) + .await + .expect("poll cancelled job"); + assert_eq!(terminal.state, CommandJobState::Cancelled); + tokio::time::sleep(Duration::from_millis(900)).await; + assert!( + !sentinel.exists(), + "shutdown cancellation left process alive" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn cancel_all_permanently_rejects_future_starts() { + let root = workspace("shutdown-reject"); + let manager = CommandJobManager::new(); + manager.cancel_all().await; + + let error = manager + .start("echo should-not-run".to_string(), root.clone(), 5_000, None) + .await + .expect_err("shutdown manager must reject new jobs"); + assert!(error.contains("shutting down")); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn start_racing_with_shutdown_cannot_escape_cancellation() { + use tokio::sync::Barrier; + + let root = workspace("shutdown-race"); + let sentinel = root.join("escaped.txt"); + let manager = CommandJobManager::new(); + let barrier = Arc::new(Barrier::new(3)); + let command = if cfg!(windows) { + "Start-Sleep -Seconds 2; Set-Content escaped.txt survived" + } else { + "sleep 2; printf survived > escaped.txt" + }; + + let starter_manager = manager.clone(); + let starter_root = root.clone(); + let starter_barrier = barrier.clone(); + let starter = tokio::spawn(async move { + starter_barrier.wait().await; + starter_manager + .start(command.to_string(), starter_root, 10_000, None) + .await + }); + + let shutdown_manager = manager.clone(); + let shutdown_barrier = barrier.clone(); + let shutdown = tokio::spawn(async move { + shutdown_barrier.wait().await; + shutdown_manager.cancel_all().await; + }); + + barrier.wait().await; + let start_result = starter.await.expect("starter task"); + shutdown.await.expect("shutdown task"); + + match start_result { + Ok(started) => { + let terminal = wait_terminal(&manager, &started.snapshot.job_id).await; + assert_eq!(terminal.state, CommandJobState::Cancelled); + } + Err(error) => assert!(error.contains("shutting down")), + } + + tokio::time::sleep(Duration::from_millis(2_200)).await; + assert!( + !sentinel.exists(), + "a start racing with shutdown escaped manager ownership" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn cleanup_prunes_expired_idempotency_keys_without_job_eviction() { + let root = workspace("dedupe-cleanup"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Write-Output done" + } else { + "printf 'done\\n'" + }; + let started = manager + .start( + command.to_string(), + root.clone(), + 5_000, + Some("expired-request-key".into()), + ) + .await + .expect("start job"); + let _ = wait_terminal(&manager, &started.snapshot.job_id).await; + + { + let mut state = manager.inner.write().await; + let entry = state + .request_jobs + .get_mut("expired-request-key") + .expect("request key exists before cleanup"); + entry.1 = Instant::now() - IDEMPOTENCY_WINDOW - StdDuration::from_secs(1); + assert!(state.jobs.contains_key(&started.snapshot.job_id)); + } + + manager.cleanup().await; + let state = manager.inner.read().await; + assert!( + !state.request_jobs.contains_key("expired-request-key"), + "expired idempotency metadata must be pruned even when no job is evicted" + ); + assert!( + state.jobs.contains_key(&started.snapshot.job_id), + "cleanup should retain the still-fresh terminal job" + ); + drop(state); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn duplicate_request_key_reuses_existing_job() { + let root = workspace("dedup"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 300" + } else { + "sleep 0.3" + }; + let first = manager + .start( + command.to_string(), + root.clone(), + 5_000, + Some("request-1".into()), + ) + .await + .expect("start first job"); + let second = manager + .start( + command.to_string(), + root.clone(), + 5_000, + Some("request-1".into()), + ) + .await + .expect("deduplicate job"); + assert!(!first.deduplicated); + assert!(second.deduplicated); + assert_eq!(first.snapshot.job_id, second.snapshot.job_id); + let _ = manager.cancel(&first.snapshot.job_id).await; + let _ = std::fs::remove_dir_all(root); + } + + #[test] + fn background_timeout_validation_covers_boundaries() { + assert_eq!( + CommandJobManager::normalize_timeout(None).expect("default timeout"), + DEFAULT_JOB_TIMEOUT_MS + ); + assert!(CommandJobManager::normalize_timeout(Some(0)).is_err()); + assert_eq!( + CommandJobManager::normalize_timeout(Some(1)).expect("minimum timeout"), + 1 + ); + assert_eq!( + CommandJobManager::normalize_timeout(Some(MAX_JOB_TIMEOUT_MS)) + .expect("maximum timeout"), + MAX_JOB_TIMEOUT_MS + ); + assert!(CommandJobManager::normalize_timeout(Some(MAX_JOB_TIMEOUT_MS + 1)).is_err()); + } + + #[tokio::test] + async fn active_job_limit_is_enforced_and_recovers_after_cancel() { + let root = workspace("capacity"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Start-Sleep -Seconds 5" + } else { + "sleep 5" + }; + let mut ids = Vec::new(); + for _ in 0..MAX_ACTIVE_JOBS { + let started = manager + .start(command.to_string(), root.clone(), 10_000, None) + .await + .expect("start capacity job"); + ids.push(started.snapshot.job_id); + } + + let overflow = manager + .start(command.to_string(), root.clone(), 10_000, None) + .await + .expect_err("ninth active job must be rejected"); + assert!(overflow.contains("too many active command jobs")); + + manager.cancel(&ids[0]).await.expect("cancel one job"); + let _ = wait_terminal(&manager, &ids[0]).await; + let replacement = manager + .start(command.to_string(), root.clone(), 10_000, None) + .await + .expect("capacity should recover after cancellation"); + ids.push(replacement.snapshot.job_id); + manager.cancel_all().await; + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn oversized_output_is_bounded_and_marks_old_cursor_truncated() { + let root = workspace("output-limit"); + let (job, _cancel_rx) = CommandJob::new("synthetic".into(), root.clone(), 5_000); + let chunk = vec![b'x'; READ_CHUNK_BYTES]; + let chunks = (MAX_OUTPUT_BYTES_PER_JOB / READ_CHUNK_BYTES) + 8; + for _ in 0..chunks { + job.append_output("stdout", &chunk).await; + } + + let snapshot = job.snapshot(0).await; + let retained = snapshot + .events + .iter() + .map(|event| event.text.len()) + .sum::(); + assert!(retained <= MAX_OUTPUT_BYTES_PER_JOB); + assert!(snapshot.output_truncated); + assert!(snapshot.events.first().is_some_and(|event| event.seq > 1)); + assert_eq!( + snapshot.next_cursor, + snapshot.events.last().expect("retained poll events").seq + ); + assert!(snapshot.has_more_output); + assert!(snapshot.next_cursor < chunks as u64); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn poll_output_is_bounded_and_cursor_drains_terminal_logs_without_gaps() { + let root = workspace("poll-output-limit"); + let (job, _cancel_rx) = CommandJob::new("synthetic".into(), root.clone(), 5_000); + let chunk = vec![b'x'; READ_CHUNK_BYTES]; + let chunks = 40usize; + for _ in 0..chunks { + job.append_output("stdout", &chunk).await; + } + job.finish(CommandJobState::Succeeded, Some(0)).await; + + let mut after = 0u64; + let mut seen = Vec::new(); + let mut polls = 0usize; + loop { + let snapshot = job.snapshot(after).await; + polls += 1; + assert_eq!(snapshot.state, CommandJobState::Succeeded); + let returned_bytes = snapshot + .events + .iter() + .map(|event| event.text.len()) + .sum::(); + assert!( + returned_bytes <= MAX_POLL_OUTPUT_BYTES, + "poll returned {returned_bytes} bytes, limit is {MAX_POLL_OUTPUT_BYTES}" + ); + for event in &snapshot.events { + assert_eq!(event.seq, after + 1, "cursor skipped or repeated an event"); + after = event.seq; + seen.push(event.seq); + } + assert_eq!(snapshot.next_cursor, after); + if !snapshot.has_more_output { + break; + } + assert!( + !snapshot.events.is_empty(), + "hasMoreOutput must make progress" + ); + } + + assert!( + polls > 1, + "test must exercise multiple bounded poll responses" + ); + assert_eq!(seen.len(), chunks); + assert_eq!(seen, (1..=chunks as u64).collect::>()); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn unknown_job_ids_are_rejected_for_poll_and_cancel() { + let manager = CommandJobManager::new(); + let poll_error = manager + .poll("definitely-not-a-job", 0, 0) + .await + .expect_err("unknown poll must fail"); + assert!(poll_error.contains("unknown or expired command job")); + let cancel_error = manager + .cancel("definitely-not-a-job") + .await + .expect_err("unknown cancel must fail"); + assert!(cancel_error.contains("unknown or expired command job")); + } + + #[tokio::test] + async fn cancelling_terminal_job_is_idempotent() { + let root = workspace("terminal-cancel"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Write-Output done" + } else { + "printf 'done\\n'" + }; + let started = manager + .start(command.to_string(), root.clone(), 5_000, None) + .await + .expect("start job"); + let terminal = wait_terminal(&manager, &started.snapshot.job_id).await; + assert_eq!(terminal.state, CommandJobState::Succeeded); + let cancelled = manager + .cancel(&started.snapshot.job_id) + .await + .expect("cancel terminal job"); + assert_eq!(cancelled.state, CommandJobState::Succeeded); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn poll_wait_returns_near_requested_deadline_when_nothing_changes() { + let root = workspace("poll-wait"); + let manager = CommandJobManager::new(); + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 500" + } else { + "sleep 0.5" + }; + let started = manager + .start(command.to_string(), root.clone(), 5_000, None) + .await + .expect("start job"); + let started_wait = Instant::now(); + let snapshot = manager + .poll(&started.snapshot.job_id, 0, 100) + .await + .expect("poll job"); + let elapsed = started_wait.elapsed(); + assert_eq!(snapshot.state, CommandJobState::Running); + assert!(snapshot.events.is_empty()); + assert!( + elapsed >= StdDuration::from_millis(70), + "poll returned too early: {elapsed:?}" + ); + assert!( + elapsed < StdDuration::from_millis(400), + "poll waited too long: {elapsed:?}" + ); + manager.cancel_all().await; + let _ = std::fs::remove_dir_all(root); + } +} diff --git a/src/main.rs b/src/main.rs index c5d8199..899f377 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,13 @@ mod binagotchy_gen; mod browser; mod command; +mod command_jobs; mod devtools; mod macos_terminal; mod mascot; mod mcp; mod ngrok; +mod process_runner; mod server; mod state; mod theme; @@ -914,6 +916,8 @@ async fn main() -> Result<(), Box> { stdout().execute(LeaveAlternateScreen)?; // Cleanup after the TUI is gone so quit never appears frozen on screen. + let command_jobs = { state.lock().await.command_jobs.clone() }; + command_jobs.cancel_all().await; { let mut app = state.lock().await; if let Some(handle) = app.server_handle.take() { @@ -3000,11 +3004,17 @@ async fn start_services( None }; - let mcp_path = { + let (mcp_path, command_jobs) = { let app = state.lock().await; - app.mcp_path() + (app.mcp_path(), app.command_jobs.clone()) }; - let router = server::router(state.clone(), devtools_bridge.clone(), mcp_path, ui_events); + let router = server::router( + state.clone(), + devtools_bridge.clone(), + command_jobs, + mcp_path, + ui_events, + ); let listener = match tokio::net::TcpListener::bind(format!("0.0.0.0:{port}")).await { Ok(l) => l, Err(e) => { diff --git a/src/mcp.rs b/src/mcp.rs index 13950a1..17825fb 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -9,6 +9,9 @@ use tiktoken_rs::o200k_base_singleton; use tokio::sync::Mutex; use crate::command; +use crate::command_jobs::{ + CommandJobManager, CommandJobSnapshot, CommandJobState, MAX_JOB_TIMEOUT_MS, MAX_POLL_WAIT_MS, +}; use crate::devtools::DevtoolsBridge; use crate::mascot; use crate::state::{ @@ -147,6 +150,7 @@ pub async fn handle_request( mode: Mode, tool_mode: ToolMode, set_catdesk_as_co_author: bool, + command_jobs: &CommandJobManager, devtools: &Option>>, ) -> Option { match req.method.as_str() { @@ -182,6 +186,7 @@ pub async fn handle_request( mode, tool_mode, set_catdesk_as_co_author, + command_jobs, devtools, ) .await, @@ -397,6 +402,43 @@ fn local_tool_output_schema(name: &str) -> Option { properties.insert("path".to_string(), json!({ "type": "string" })); properties.insert("recursive".to_string(), json!({ "type": "boolean" })); } + "start_command" | "poll_command" | "cancel_command" => { + for field in ["jobId", "command", "cwd", "state"] { + properties.insert(field.to_string(), json!({ "type": "string" })); + } + for field in ["elapsedMs", "nextCursor", "timeoutMs"] { + properties.insert( + field.to_string(), + json!({ "type": "integer", "minimum": 0 }), + ); + } + properties.insert( + "exitCode".to_string(), + json!({ "type": ["integer", "null"] }), + ); + properties.insert( + "commandSuccess".to_string(), + json!({ "type": ["boolean", "null"] }), + ); + properties.insert("hasMoreOutput".to_string(), json!({ "type": "boolean" })); + properties.insert("outputTruncated".to_string(), json!({ "type": "boolean" })); + properties.insert("deduplicated".to_string(), json!({ "type": "boolean" })); + properties.insert( + "events".to_string(), + json!({ + "type": "array", + "items": { + "type": "object", + "properties": { + "seq": { "type": "integer", "minimum": 1 }, + "stream": { "type": "string", "enum": ["stdout", "stderr"] }, + "text": { "type": "string" } + }, + "required": ["seq", "stream", "text"] + } + }), + ); + } "run_command" => { for field in [ "command", @@ -427,11 +469,18 @@ fn local_tool_output_schema(name: &str) -> Option { json!({ "type": "integer", "minimum": 0 }), ); } + properties.insert( + "exitCode".to_string(), + json!({ "type": ["integer", "null"] }), + ); for field in [ "destinationOperandWasDirectory", "overwrite", "skipped", "listTruncated", + "timedOut", + "stdoutTruncated", + "stderrTruncated", ] { properties.insert(field.to_string(), json!({ "type": "boolean" })); } @@ -502,12 +551,55 @@ async fn handle_tools_list( "properties": { "command": { "type": "string", "description": "The shell command to execute" }, "cwd": { "type": "string", "description": "Working directory relative to workspace root or absolute path within it" }, - "timeout": { "type": "number", "description": "Timeout in milliseconds. Clamped to 120000." } + "timeout": { "type": "integer", "minimum": 1, "maximum": 120000, "description": "Timeout in milliseconds for short commands. Maximum 120000; use start_command for long-running work." } }, "required": ["command"] }, "annotations": { "readOnlyHint": false, "openWorldHint": true, "destructiveHint": true } })); + tools.push(json!({ + "name": "start_command", + "title": "Start command", + "description": "Start a long-running shell command inside the workspace and return a job ID immediately. Prefer this for builds, compilation, dependency installation, long test suites, and development servers instead of keeping run_command open.", + "inputSchema": { + "type": "object", + "properties": { + "command": { "type": "string", "description": "The shell command to start" }, + "cwd": { "type": "string", "description": "Working directory relative to workspace root or absolute path within it" }, + "timeout": { "type": "integer", "minimum": 1, "maximum": MAX_JOB_TIMEOUT_MS, "description": "Maximum command runtime in milliseconds. Defaults to 30 minutes; maximum is 24 hours." } + }, + "required": ["command"] + }, + "annotations": { "readOnlyHint": false, "openWorldHint": true, "destructiveHint": true } + })); + tools.push(json!({ + "name": "poll_command", + "title": "Poll command", + "description": "Read incremental output and current status from a command previously started with start_command. Pass the returned nextCursor as after on the next poll so output is not repeated. If hasMoreOutput is true, poll again even if the job is already terminal so the remaining buffered output can be drained.", + "inputSchema": { + "type": "object", + "properties": { + "job_id": { "type": "string", "description": "Opaque command job ID returned by start_command" }, + "after": { "type": "integer", "minimum": 0, "description": "Return only output events after this cursor (default 0)" }, + "wait_ms": { "type": "integer", "minimum": 0, "maximum": MAX_POLL_WAIT_MS, "description": "Wait briefly for new output or completion before returning (maximum 30000 ms)" } + }, + "required": ["job_id"] + }, + "annotations": { "readOnlyHint": false, "openWorldHint": false, "destructiveHint": false } + })); + tools.push(json!({ + "name": "cancel_command", + "title": "Cancel command", + "description": "Cancel a command started with start_command and terminate its complete child process tree.", + "inputSchema": { + "type": "object", + "properties": { + "job_id": { "type": "string", "description": "Opaque command job ID returned by start_command" } + }, + "required": ["job_id"] + }, + "annotations": { "readOnlyHint": false, "openWorldHint": false, "destructiveHint": true } + })); } tools.push(json!({ @@ -637,6 +729,7 @@ async fn handle_tools_call( mode: Mode, tool_mode: ToolMode, set_catdesk_as_co_author: bool, + command_jobs: &CommandJobManager, devtools: &Option>>, ) -> JsonRpcResponse { let params = &req.params; @@ -652,9 +745,28 @@ async fn handle_tools_call( let mut response = { // Local computer tools if mode.computer_enabled() { - if tool_name == "run_command" { + if matches!( + tool_name.as_str(), + "run_command" | "start_command" | "poll_command" | "cancel_command" + ) { if tool_mode.run_command_enabled() { - handle_run_command(req, workspace_root, set_catdesk_as_co_author).await + match tool_name.as_str() { + "run_command" => { + handle_run_command(req, workspace_root, set_catdesk_as_co_author).await + } + "start_command" => { + handle_start_command( + req, + workspace_root, + set_catdesk_as_co_author, + command_jobs, + ) + .await + } + "poll_command" => handle_poll_command(req, command_jobs).await, + "cancel_command" => handle_cancel_command(req, command_jobs).await, + _ => unreachable!(), + } } else if tool_mode.read_only() { read_only_blocked_response(req, &tool_name) } else { @@ -794,6 +906,208 @@ async fn forward_to_devtools( } } +fn command_job_output_text(snapshot: &CommandJobSnapshot) -> String { + if snapshot.events.is_empty() { + return match snapshot.state { + CommandJobState::Running => "(no new output; command is still running)".to_string(), + _ => "(no new output)".to_string(), + }; + } + let mut output = String::new(); + for event in &snapshot.events { + if event.stream == "stderr" { + output.push_str("[stderr] "); + } + output.push_str(&event.text); + if !event.text.ends_with('\n') { + output.push('\n'); + } + } + if snapshot.has_more_output { + output.push_str("[more buffered output available; poll again with nextCursor]\n"); + } + output +} + +fn command_job_structured(tool_name: &str, snapshot: &CommandJobSnapshot) -> Value { + let command_success = match snapshot.state { + CommandJobState::Succeeded => Some(true), + CommandJobState::Failed | CommandJobState::Cancelled | CommandJobState::TimedOut => { + Some(false) + } + CommandJobState::Running => None, + }; + json!({ + "toolName": tool_name, + "jobId": snapshot.job_id, + "command": snapshot.command, + "cwd": snapshot.cwd, + "state": snapshot.state.as_str(), + "elapsedMs": snapshot.elapsed_ms, + "exitCode": snapshot.exit_code, + "events": snapshot.events, + "nextCursor": snapshot.next_cursor, + "hasMoreOutput": snapshot.has_more_output, + "outputTruncated": snapshot.output_truncated, + "timeoutMs": snapshot.timeout_ms, + "commandSuccess": command_success, + "success": true, + }) +} + +async fn handle_start_command( + req: &JsonRpcRequest, + workspace_root: &str, + set_catdesk_as_co_author: bool, + command_jobs: &CommandJobManager, +) -> JsonRpcResponse { + let arguments = tool_arguments(req); + let command_text = match required_string_argument(&arguments, "command") { + Ok(value) => value, + Err(error) => return tool_error_response(req, error), + }; + if command::contains_catdesk_co_author_marker(command_text) { + let message = if set_catdesk_as_co_author { + "Rewrite the commit message normally and remove \"Co-Authored-By: CatDesk\". CatDesk will add that trailer automatically." + } else { + "Do not include \"Co-Authored-By: CatDesk\" in the commit message. The user does not want that attribution." + }; + return tool_error_response(req, message.into()); + } + let cwd_input = match optional_string_argument(&arguments, "cwd") { + Ok(value) => value, + Err(error) => return tool_error_response(req, error), + }; + let cwd = match command::resolve_workspace_path(workspace_root, cwd_input) { + Ok(path) => path, + Err(error) => { + return tool_error_response( + req, + format!("code: PATH_OUTSIDE_WORKSPACE\nmessage: {error}"), + ); + } + }; + let requested_timeout = match arguments.get("timeout") { + Some(value) => match value.as_u64() { + Some(value) => Some(value), + None => { + return tool_error_response( + req, + "Parameter timeout must be a positive integer".into(), + ); + } + }, + None => None, + }; + let timeout_ms = match CommandJobManager::normalize_timeout(requested_timeout) { + Ok(value) => value, + Err(error) => return tool_error_response(req, error), + }; + let effective_command = + if set_catdesk_as_co_author && command::command_contains_git_commit(command_text) { + command::inject_catdesk_co_author_trailer(command_text) + } else { + command_text.to_string() + }; + let request_key = req.id.as_ref().map(|id| { + let mut hasher = DefaultHasher::new(); + effective_command.hash(&mut hasher); + cwd.hash(&mut hasher); + timeout_ms.hash(&mut hasher); + format!("start_command:{id}:{:016x}", hasher.finish()) + }); + match command_jobs + .start(effective_command, cwd, timeout_ms, request_key) + .await + { + Ok(started) => { + let mut structured = command_job_structured("start_command", &started.snapshot); + if let Some(object) = structured.as_object_mut() { + object.insert("deduplicated".to_string(), json!(started.deduplicated)); + } + let text = if started.deduplicated { + format!("Command job already exists: {}", started.snapshot.job_id) + } else { + format!("Started command job: {}", started.snapshot.job_id) + }; + tool_success_response_with_structured(req, text, structured) + } + Err(error) => tool_error_response(req, error), + } +} + +async fn handle_poll_command( + req: &JsonRpcRequest, + command_jobs: &CommandJobManager, +) -> JsonRpcResponse { + let arguments = tool_arguments(req); + let job_id = match required_string_argument(&arguments, "job_id") { + Ok(value) => value, + Err(error) => return tool_error_response(req, error), + }; + let after = match arguments.get("after") { + Some(value) => match value.as_u64() { + Some(value) => value, + None => { + return tool_error_response( + req, + "Parameter after must be a non-negative integer".into(), + ); + } + }, + None => 0, + }; + let wait_ms = match arguments.get("wait_ms") { + Some(value) => match value.as_u64() { + Some(value) if value <= MAX_POLL_WAIT_MS => value, + Some(_) => { + return tool_error_response( + req, + format!("wait_ms must be at most {MAX_POLL_WAIT_MS}"), + ); + } + None => { + return tool_error_response( + req, + "Parameter wait_ms must be a non-negative integer".into(), + ); + } + }, + None => 0, + }; + match command_jobs.poll(job_id, after, wait_ms).await { + Ok(snapshot) => { + let text = command_job_output_text(&snapshot); + let structured = command_job_structured("poll_command", &snapshot); + tool_success_response_with_structured(req, text, structured) + } + Err(error) => tool_error_response(req, error), + } +} + +async fn handle_cancel_command( + req: &JsonRpcRequest, + command_jobs: &CommandJobManager, +) -> JsonRpcResponse { + let arguments = tool_arguments(req); + let job_id = match required_string_argument(&arguments, "job_id") { + Ok(value) => value, + Err(error) => return tool_error_response(req, error), + }; + match command_jobs.cancel(job_id).await { + Ok(snapshot) => { + let text = format!( + "Command job {} is {}", + snapshot.job_id, + snapshot.state.as_str() + ); + let structured = command_job_structured("cancel_command", &snapshot); + tool_success_response_with_structured(req, text, structured) + } + Err(error) => tool_error_response(req, error), + } +} + async fn handle_run_command( req: &JsonRpcRequest, workspace_root: &str, @@ -810,6 +1124,20 @@ async fn handle_run_command( let cwd_input = arguments.get("cwd").and_then(|v| v.as_str()); let timeout_ms = arguments.get("timeout").and_then(|v| v.as_u64()); + if let Some(timeout_ms) = timeout_ms { + if timeout_ms == 0 { + return tool_error_response(req, "timeout must be at least 1 ms".into()); + } + if timeout_ms > command::MAX_TIMEOUT_MS { + return tool_error_response( + req, + format!( + "run_command supports at most {} ms. Use start_command for builds, compilation, dependency installation, long test suites, development servers, or other long-running commands.", + command::MAX_TIMEOUT_MS + ), + ); + } + } if command::contains_catdesk_co_author_marker(cmd) { let message = if set_catdesk_as_co_author { @@ -888,7 +1216,11 @@ async fn handle_run_command( "stdout": result.stdout, "stderr": result.stderr, "success": result.success, + "exitCode": result.exit_code, "elapsedMs": result.elapsed_ms, + "timedOut": result.timed_out, + "stdoutTruncated": result.stdout_truncated, + "stderrTruncated": result.stderr_truncated, }); if result.success { @@ -1345,7 +1677,19 @@ Always specify the branch explicitly when using `git push`."# if mode.computer_enabled() && tool_mode.run_command_enabled() { lines.push( - "Use run_command only as a last resort when the available dedicated tools cannot complete the operation." + "Use run_command only as a last resort when the available dedicated tools cannot complete the operation, and keep it for short commands that should finish quickly." + .to_string(), + ); + lines.push( + "For builds, compilation, dependency installation, long-running test suites, development servers, or commands that may take more than about one minute, use start_command instead of keeping run_command open." + .to_string(), + ); + lines.push( + "Use poll_command to read incremental output from a background command. Pass the returned nextCursor as after on the next poll so output is not repeated, and use wait_ms when waiting briefly for new progress. If hasMoreOutput is true, keep polling even after the command reaches a terminal state so all buffered output can be drained." + .to_string(), + ); + lines.push( + "Use cancel_command when a background command is no longer needed. Do not repeatedly start the same build or server while an existing command job is still running." .to_string(), ); } @@ -1603,7 +1947,16 @@ fn attach_tool_call_count(result: &mut Value, tool_call_count: u64) { fn tool_descriptor_should_attach_widget(name: &str) -> bool { matches!( name, - "run_command" | "catdesk_instruction" | "search" | "read" | "write" | "edit" | "delete" + "run_command" + | "start_command" + | "poll_command" + | "cancel_command" + | "catdesk_instruction" + | "search" + | "read" + | "write" + | "edit" + | "delete" ) } @@ -1829,11 +2182,13 @@ fn current_token_stats_layout() -> TokenStatsLayout { .unwrap_or_default() } +#[cfg(test)] +fn current_show_detail_mode() -> ShowDetailMode { + ShowDetailMode::Expanded +} + +#[cfg(not(test))] fn current_show_detail_mode() -> ShowDetailMode { - #[cfg(test)] - { - return ShowDetailMode::Expanded; - } crate::state::load_app_config() .map(|config| config.show_detail_mode) .unwrap_or_default() @@ -2011,6 +2366,75 @@ fn build_run_command_widget_payload( Some(Value::Object(payload)) } +fn build_command_job_widget_payload(result: &Value, tool_name: &str) -> Option { + let structured = result_structured_content(result)?; + let command = structured.get("command")?.clone(); + let state = structured.get("state")?.as_str()?; + let (title, widget_state) = match state { + "running" => ( + if tool_name == "start_command" { + "Command Started" + } else { + "Command Running" + }, + "waiting", + ), + "succeeded" => ("Command Complete", "done"), + "cancelled" => ("Command Cancelled", "done"), + "failed" => ("Command Failed", "failed"), + "timed_out" => ("Command Timed Out", "failed"), + _ => ("Command Job", "waiting"), + }; + let mut output = String::new(); + if let Some(events) = structured.get("events").and_then(Value::as_array) { + for event in events { + let stream = event + .get("stream") + .and_then(Value::as_str) + .unwrap_or("stdout"); + let text = event + .get("text") + .and_then(Value::as_str) + .unwrap_or_default(); + if stream == "stderr" { + output.push_str("[stderr] "); + } + output.push_str(text); + if !text.ends_with('\n') { + output.push('\n'); + } + } + } + if output.is_empty() { + output = format!( + "job {} ยท {}", + structured + .get("jobId") + .and_then(Value::as_str) + .unwrap_or("?"), + state + ); + } + if structured.get("outputTruncated").and_then(Value::as_bool) == Some(true) { + output.push_str("\n[older command output was truncated]\n"); + } + if structured.get("hasMoreOutput").and_then(Value::as_bool) == Some(true) { + output.push_str("\n[more buffered output available; poll again]\n"); + } + let mut payload = base_widget_payload("tool_call", title, widget_state, Some(tool_name)); + payload.insert("command".to_string(), command); + payload.insert( + "output".to_string(), + json!(truncate_for_widget(&output, MAX_COMMAND_OUTPUT_CHARS)), + ); + if let Some(elapsed) = structured.get("elapsedMs") { + payload.insert("elapsedMs".to_string(), elapsed.clone()); + } + payload.insert("changedFiles".to_string(), json!([])); + payload.insert("hasChanges".to_string(), json!(false)); + Some(Value::Object(payload)) +} + fn build_generic_widget_payload( req: &JsonRpcRequest, result: &Value, @@ -2143,6 +2567,19 @@ fn build_auto_widget_payload( "Failed to build run_command widget payload from structuredContent.".into(), ), }, + "start_command" | "poll_command" | "cancel_command" => { + match build_command_job_widget_payload(result, &tool_name) { + Some(payload) => payload, + None if is_error => { + build_generic_widget_payload(req, result, widget_context, is_error) + } + None => build_widget_payload_error( + req, + widget_context, + format!("Failed to build {tool_name} widget payload from structuredContent."), + ), + } + } _ => build_generic_widget_payload(req, result, widget_context, is_error), } } @@ -2697,7 +3134,16 @@ fn diff_changed_files(before: &WatchedSnapshot, after: &WatchedSnapshot) -> Vec< } fn is_local_destructive_tool(tool_name: &str) -> bool { - matches!(tool_name, "run_command" | "write" | "edit" | "delete") + matches!( + tool_name, + "run_command" + | "start_command" + | "poll_command" + | "cancel_command" + | "write" + | "edit" + | "delete" + ) } fn tool_is_read_only(tool: &Value) -> bool { @@ -3072,6 +3518,477 @@ mod tests { assert_eq!(DEVTOOLS_PROTOCOL_VERSION, "2025-03-26"); } + #[tokio::test] + async fn command_job_tools_start_poll_and_report_terminal_success() { + let workspace_root = + std::env::temp_dir().join(format!("catdesk-mcp-command-job-{}", Uuid::new_v4())); + std::fs::create_dir_all(&workspace_root).expect("create workspace"); + let workspace_root_str = workspace_root.to_string_lossy().into_owned(); + let command_jobs = CommandJobManager::new(); + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 150; Write-Output job-done" + } else { + "sleep 0.15; printf 'job-done\\n'" + }; + + let start_req = tool_call_request( + "start_command", + json!({ "command": command, "timeout": 5_000 }), + ); + let start_response = handle_tools_call( + &start_req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::MultiTools, + false, + &command_jobs, + &None, + ) + .await; + assert_no_text_content(&start_response); + let start_structured = start_response + .result + .as_ref() + .and_then(|result| result.get("structuredContent")) + .expect("missing start structured content"); + let job_id = start_structured + .get("jobId") + .and_then(Value::as_str) + .expect("missing job id") + .to_string(); + assert_eq!( + start_structured.get("state").and_then(Value::as_str), + Some("running") + ); + assert_eq!( + start_response + .result + .as_ref() + .and_then(|result| result.get("_meta")) + .and_then(|meta| meta.get(WIDGET_PAYLOAD_META_KEY)) + .and_then(|payload| payload.get("toolName")) + .and_then(Value::as_str), + Some("start_command") + ); + + let mut terminal = None; + let mut cursor = 0; + let mut seen_output = String::new(); + for _ in 0..20 { + let poll_req = tool_call_request( + "poll_command", + json!({ "job_id": job_id, "after": cursor, "wait_ms": 250 }), + ); + let response = handle_tools_call( + &poll_req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::MultiTools, + false, + &command_jobs, + &None, + ) + .await; + let structured = response + .result + .as_ref() + .and_then(|result| result.get("structuredContent")) + .expect("missing poll structured content"); + if let Some(events) = structured.get("events").and_then(Value::as_array) { + for event in events { + if let Some(text) = event.get("text").and_then(Value::as_str) { + seen_output.push_str(text); + } + } + } + cursor = structured + .get("nextCursor") + .and_then(Value::as_u64) + .unwrap_or(cursor); + if structured.get("state").and_then(Value::as_str) == Some("succeeded") + && structured.get("hasMoreOutput").and_then(Value::as_bool) != Some(true) + { + terminal = Some(response); + break; + } + } + let terminal = terminal.expect("job did not reach succeeded state"); + assert!( + terminal + .result + .as_ref() + .and_then(|result| result.get("isError")) + .is_none(), + "successful command polling must not be an MCP tool error" + ); + let structured = terminal + .result + .as_ref() + .and_then(|result| result.get("structuredContent")) + .expect("missing terminal structured content"); + assert_eq!( + structured.get("commandSuccess").and_then(Value::as_bool), + Some(true) + ); + assert!(seen_output.contains("job-done")); + + let _ = std::fs::remove_dir_all(workspace_root); + } + + #[tokio::test] + async fn reused_json_rpc_id_with_different_start_arguments_creates_distinct_jobs() { + let workspace_root = + std::env::temp_dir().join(format!("catdesk-mcp-id-reuse-{}", Uuid::new_v4())); + std::fs::create_dir_all(&workspace_root).expect("create workspace"); + let workspace_root_str = workspace_root.to_string_lossy().into_owned(); + let command_jobs = CommandJobManager::new(); + let first_command = if cfg!(windows) { + "Start-Sleep -Milliseconds 500" + } else { + "sleep 0.5" + }; + let second_command = if cfg!(windows) { + "Start-Sleep -Milliseconds 600" + } else { + "sleep 0.6" + }; + + // tool_call_request deliberately reuses the same JSON-RPC id. Stateless + // clients are allowed to do this across independent calls. + let first_req = tool_call_request("start_command", json!({ "command": first_command })); + let second_req = tool_call_request("start_command", json!({ "command": second_command })); + let first = handle_tools_call( + &first_req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::MultiTools, + false, + &command_jobs, + &None, + ) + .await; + let second = handle_tools_call( + &second_req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::MultiTools, + false, + &command_jobs, + &None, + ) + .await; + + let job_id = |response: &JsonRpcResponse| { + response + .result + .as_ref() + .and_then(|result| result.get("structuredContent")) + .and_then(|structured| structured.get("jobId")) + .and_then(Value::as_str) + .expect("missing job id") + .to_string() + }; + assert_ne!(job_id(&first), job_id(&second)); + command_jobs.cancel_all().await; + let _ = std::fs::remove_dir_all(workspace_root); + } + + #[test] + fn command_job_widget_state_matrix_preserves_command_ui_contract() { + let cases = [ + ("start_command", "running", "Command Started", "waiting"), + ("poll_command", "running", "Command Running", "waiting"), + ("poll_command", "succeeded", "Command Complete", "done"), + ("poll_command", "failed", "Command Failed", "failed"), + ("cancel_command", "cancelled", "Command Cancelled", "done"), + ("poll_command", "timed_out", "Command Timed Out", "failed"), + ]; + + for (tool_name, state, expected_title, expected_widget_state) in cases { + let result = json!({ + "structuredContent": { + "toolName": tool_name, + "jobId": "job-123", + "command": "cargo build", + "cwd": "E:/CatDesk", + "state": state, + "elapsedMs": 123, + "exitCode": null, + "events": [], + "nextCursor": 0, + "outputTruncated": false, + "timeoutMs": 5000, + "commandSuccess": null, + "success": true + } + }); + let payload = build_command_job_widget_payload(&result, tool_name) + .unwrap_or_else(|| panic!("missing widget payload for {tool_name}/{state}")); + assert_eq!( + payload.get("toolName").and_then(Value::as_str), + Some(tool_name) + ); + assert_eq!( + payload.get("title").and_then(Value::as_str), + Some(expected_title) + ); + assert_eq!( + payload.get("state").and_then(Value::as_str), + Some(expected_widget_state) + ); + assert_eq!( + payload.get("command").and_then(Value::as_str), + Some("cargo build") + ); + assert_eq!(payload.get("elapsedMs").and_then(Value::as_u64), Some(123)); + assert_eq!( + payload.get("hasChanges").and_then(Value::as_bool), + Some(false) + ); + } + } + + #[test] + fn command_job_widget_formats_stderr_and_truncation_without_new_styles() { + let result = json!({ + "structuredContent": { + "toolName": "poll_command", + "jobId": "job-123", + "command": "cargo build", + "cwd": "E:/CatDesk", + "state": "failed", + "elapsedMs": 456, + "exitCode": 1, + "events": [ + {"seq": 4, "stream": "stdout", "text": "compiling\n"}, + {"seq": 5, "stream": "stderr", "text": "error: nope\n"} + ], + "nextCursor": 5, + "hasMoreOutput": true, + "outputTruncated": true, + "timeoutMs": 5000, + "commandSuccess": false, + "success": true + } + }); + let payload = build_command_job_widget_payload(&result, "poll_command") + .expect("command job widget payload"); + let output = payload + .get("output") + .and_then(Value::as_str) + .expect("missing widget output"); + assert!(output.contains("compiling")); + assert!(output.contains("[stderr] error: nope")); + assert!(output.contains("[older command output was truncated]")); + assert!(output.contains("[more buffered output available; poll again]")); + assert_eq!( + payload.get("title").and_then(Value::as_str), + Some("Command Failed") + ); + assert_eq!(payload.get("state").and_then(Value::as_str), Some("failed")); + } + + #[test] + fn original_run_command_widget_shape_is_unchanged_by_new_runtime_metadata() { + let req = tool_call_request("run_command", json!({ "command": "cargo check" })); + let raw = json!({ + "content": [], + "structuredContent": { + "toolName": "run_command", + "command": "cargo check", + "cwd": "E:/CatDesk", + "stdout": "Finished dev profile\n", + "stderr": "", + "success": true, + "exitCode": 0, + "elapsedMs": 321, + "timedOut": false, + "stdoutTruncated": false, + "stderrTruncated": false + } + }); + let result = enrich_tool_result(&req, raw, None); + let payload = result + .get("_meta") + .and_then(|meta| meta.get(WIDGET_PAYLOAD_META_KEY)) + .expect("missing run_command widget payload"); + assert_eq!( + payload.get("toolName").and_then(Value::as_str), + Some("run_command") + ); + assert_eq!( + payload.get("title").and_then(Value::as_str), + Some("Command Output") + ); + assert_eq!(payload.get("state").and_then(Value::as_str), Some("done")); + assert_eq!( + payload.get("command").and_then(Value::as_str), + Some("cargo check") + ); + assert_eq!(payload.get("elapsedMs").and_then(Value::as_u64), Some(321)); + assert!(payload.get("exitCode").is_none()); + assert!(payload.get("timedOut").is_none()); + assert!(payload.get("stdoutTruncated").is_none()); + assert!(payload.get("stderrTruncated").is_none()); + } + + #[tokio::test] + async fn read_only_mode_blocks_all_command_job_calls_even_if_invoked_directly() { + let workspace_root = + std::env::temp_dir().join(format!("catdesk-mcp-command-read-only-{}", Uuid::new_v4())); + std::fs::create_dir_all(&workspace_root).expect("create workspace"); + let workspace_root_str = workspace_root.to_string_lossy().into_owned(); + let command_jobs = CommandJobManager::new(); + + for (tool_name, arguments) in [ + ("start_command", json!({"command": "echo blocked"})), + ("poll_command", json!({"job_id": "blocked"})), + ("cancel_command", json!({"job_id": "blocked"})), + ] { + let req = tool_call_request(tool_name, arguments); + let response = handle_tools_call( + &req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::ReadOnly, + false, + &command_jobs, + &None, + ) + .await; + assert_eq!( + response + .result + .as_ref() + .and_then(|result| result.get("isError")) + .and_then(Value::as_bool), + Some(true), + "{tool_name} should be blocked in read-only mode" + ); + assert!(result_text(&response).contains("disabled in read-only mode")); + } + + let _ = std::fs::remove_dir_all(workspace_root); + } + + #[tokio::test] + async fn failed_background_command_is_pollable_without_mcp_error() { + let workspace_root = + std::env::temp_dir().join(format!("catdesk-mcp-command-fail-{}", Uuid::new_v4())); + std::fs::create_dir_all(&workspace_root).expect("create workspace"); + let workspace_root_str = workspace_root.to_string_lossy().into_owned(); + let command_jobs = CommandJobManager::new(); + let start_req = tool_call_request("start_command", json!({ "command": "exit 7" })); + let start_response = handle_tools_call( + &start_req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::MultiTools, + false, + &command_jobs, + &None, + ) + .await; + let job_id = start_response + .result + .as_ref() + .and_then(|result| result.get("structuredContent")) + .and_then(|structured| structured.get("jobId")) + .and_then(Value::as_str) + .expect("missing job id") + .to_string(); + + let mut terminal = None; + for _ in 0..20 { + let poll_req = + tool_call_request("poll_command", json!({ "job_id": job_id, "wait_ms": 250 })); + let response = handle_tools_call( + &poll_req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::MultiTools, + false, + &command_jobs, + &None, + ) + .await; + let state = response + .result + .as_ref() + .and_then(|result| result.get("structuredContent")) + .and_then(|structured| structured.get("state")) + .and_then(Value::as_str); + let has_more = response + .result + .as_ref() + .and_then(|result| result.get("structuredContent")) + .and_then(|structured| structured.get("hasMoreOutput")) + .and_then(Value::as_bool) + .unwrap_or(false); + if state == Some("failed") && !has_more { + terminal = Some(response); + break; + } + } + let terminal = terminal.expect("job did not reach failed state"); + let result = terminal.result.as_ref().expect("missing result"); + assert!(result.get("isError").is_none()); + let structured = result + .get("structuredContent") + .expect("missing structured content"); + assert_eq!( + structured.get("state").and_then(Value::as_str), + Some("failed") + ); + assert_eq!( + structured.get("commandSuccess").and_then(Value::as_bool), + Some(false) + ); + assert_eq!(structured.get("exitCode").and_then(Value::as_i64), Some(7)); + + let _ = std::fs::remove_dir_all(workspace_root); + } + + #[tokio::test] + async fn run_command_rejects_long_timeout_and_points_to_start_command() { + let workspace_root = + std::env::temp_dir().join(format!("catdesk-mcp-run-timeout-{}", Uuid::new_v4())); + std::fs::create_dir_all(&workspace_root).expect("create workspace"); + let workspace_root_str = workspace_root.to_string_lossy().into_owned(); + let req = tool_call_request( + "run_command", + json!({ "command": "echo short", "timeout": command::MAX_TIMEOUT_MS + 1 }), + ); + let response = handle_tools_call( + &req, + &workspace_root_str, + 1, + Mode::Both, + ToolMode::MultiTools, + false, + &CommandJobManager::new(), + &None, + ) + .await; + assert_eq!( + response + .result + .as_ref() + .and_then(|result| result.get("isError")) + .and_then(Value::as_bool), + Some(true) + ); + assert!(result_text(&response).contains("Use start_command")); + let _ = std::fs::remove_dir_all(workspace_root); + } + #[tokio::test] async fn multi_tools_list_exposes_run_command_mv_without_move_path_tool() { let req = JsonRpcRequest { @@ -3096,6 +4013,9 @@ mod tests { names, vec![ "run_command", + "start_command", + "poll_command", + "cancel_command", "catdesk_instruction", "read", "search", @@ -3302,6 +4222,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3344,6 +4265,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3376,6 +4298,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3426,6 +4349,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3525,6 +4449,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3597,6 +4522,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3676,6 +4602,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3725,6 +4652,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3800,6 +4728,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3873,6 +4802,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -3956,6 +4886,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -4018,6 +4949,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -4054,6 +4986,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; @@ -4089,6 +5022,7 @@ mod tests { Mode::Both, ToolMode::MultiTools, false, + &CommandJobManager::new(), &None, ) .await; diff --git a/src/process_runner.rs b/src/process_runner.rs new file mode 100644 index 0000000..ea1314e --- /dev/null +++ b/src/process_runner.rs @@ -0,0 +1,594 @@ +use std::io; +use std::path::Path; +use std::process::Stdio; +use std::time::Instant; + +use tokio::io::{AsyncRead, AsyncReadExt}; +use tokio::process::{Child, ChildStderr, ChildStdout, Command}; +use tokio::time::{Duration, timeout}; + +const READ_CHUNK_BYTES: usize = 8 * 1024; + +#[derive(Debug)] +pub struct ProcessRunResult { + pub stdout: String, + pub stderr: String, + pub success: bool, + pub exit_code: Option, + pub elapsed_ms: u64, + pub timed_out: bool, + pub stdout_truncated: bool, + pub stderr_truncated: bool, +} + +/// A spawned shell process owned by CatDesk. +/// +/// Dropping this value is intentionally destructive: if the command is still +/// alive, CatDesk terminates the process tree. This is what keeps a cancelled +/// MCP request from leaving a compiler or build process behind. +pub struct SpawnedProcess { + child: Child, + stdout: Option, + stderr: Option, + tree: ProcessTreeGuard, +} + +impl SpawnedProcess { + pub fn take_stdout(&mut self) -> Option { + self.stdout.take() + } + + pub fn take_stderr(&mut self) -> Option { + self.stderr.take() + } + + pub async fn wait(&mut self) -> io::Result { + self.child.wait().await + } + + /// Terminate the root process and all descendants owned by this command. + pub fn terminate_tree(&mut self) { + self.tree.terminate(); + // `taskkill /T` / process-group termination should already include the + // root, but keep Tokio's direct kill as a best-effort fallback. + let _ = self.child.start_kill(); + } + + /// Finalize ownership after the root process exits. Any descendants still + /// alive at that point are terminated so a command cannot silently detach + /// work that outlives its CatDesk job. + pub fn disarm(&mut self) { + self.tree.disarm(); + } +} + +impl Drop for SpawnedProcess { + fn drop(&mut self) { + if self.tree.is_armed() { + self.tree.terminate(); + let _ = self.child.start_kill(); + } + } +} + +#[derive(Debug)] +struct ProcessTreeGuard { + pid: u32, + armed: bool, + #[cfg(windows)] + job_handle: Option, +} + +impl ProcessTreeGuard { + fn new(pid: u32) -> Self { + Self { + pid, + armed: true, + #[cfg(windows)] + job_handle: create_windows_job_for_process(pid), + } + } + + fn is_armed(&self) -> bool { + self.armed + } + + fn disarm(&mut self) { + if !self.armed { + return; + } + #[cfg(windows)] + { + if self.job_handle.is_some() { + close_windows_job(&mut self.job_handle); + } else { + // Best effort when Job Object assignment was unavailable. + terminate_process_tree(self.pid); + } + } + #[cfg(not(windows))] + terminate_process_tree(self.pid); + self.armed = false; + } + + fn terminate(&mut self) { + if !self.armed { + return; + } + #[cfg(windows)] + { + if !terminate_windows_job(&mut self.job_handle) { + terminate_process_tree(self.pid); + } + } + #[cfg(not(windows))] + terminate_process_tree(self.pid); + self.armed = false; + } +} + +impl Drop for ProcessTreeGuard { + fn drop(&mut self) { + self.terminate(); + } +} + +#[cfg(windows)] +fn create_windows_job_for_process(pid: u32) -> Option { + use std::ffi::c_void; + use std::mem::{size_of, zeroed}; + use windows_sys::Win32::Foundation::CloseHandle; + use windows_sys::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, + SetInformationJobObject, + }; + use windows_sys::Win32::System::Threading::{ + OpenProcess, PROCESS_SET_QUOTA, PROCESS_TERMINATE, + }; + + unsafe { + let job = CreateJobObjectW(std::ptr::null(), std::ptr::null()); + if job.is_null() { + return None; + } + + let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = zeroed(); + info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + if SetInformationJobObject( + job, + JobObjectExtendedLimitInformation, + &info as *const _ as *const c_void, + size_of::() as u32, + ) == 0 + { + CloseHandle(job); + return None; + } + + let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, 0, pid); + if process.is_null() { + CloseHandle(job); + return None; + } + let assigned = AssignProcessToJobObject(job, process) != 0; + CloseHandle(process); + if !assigned { + CloseHandle(job); + return None; + } + + Some(job as usize) + } +} + +#[cfg(windows)] +fn close_windows_job(job_handle: &mut Option) { + use windows_sys::Win32::Foundation::{CloseHandle, HANDLE}; + if let Some(raw) = job_handle.take() { + unsafe { + CloseHandle(raw as HANDLE); + } + } +} + +#[cfg(windows)] +fn terminate_windows_job(job_handle: &mut Option) -> bool { + use windows_sys::Win32::Foundation::{CloseHandle, HANDLE}; + use windows_sys::Win32::System::JobObjects::TerminateJobObject; + let Some(raw) = job_handle.take() else { + return false; + }; + let handle = raw as HANDLE; + unsafe { + let terminated = TerminateJobObject(handle, 1) != 0; + CloseHandle(handle); + terminated + } +} + +#[cfg(windows)] +fn terminate_process_tree(pid: u32) { + // `/T` includes descendants and `/F` makes cancellation deterministic. + // Use the executable directly rather than a shell command so the PID never + // passes through shell parsing. + let _ = std::process::Command::new("taskkill") + .args(["/PID", &pid.to_string(), "/T", "/F"]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); +} + +#[cfg(unix)] +fn terminate_process_tree(pid: u32) { + let pgid = match i32::try_from(pid) { + Ok(value) => value, + Err(_) => return, + }; + // The shell is placed in its own process group at spawn time. A negative + // PID targets the complete process group, including compiler descendants. + unsafe { + let _ = libc::kill(-pgid, libc::SIGKILL); + } +} + +#[cfg(not(any(windows, unix)))] +fn terminate_process_tree(_pid: u32) {} + +fn shell_command(command: &str) -> Command { + #[cfg(windows)] + { + let mut shell = Command::new("powershell.exe"); + shell + .arg("-NoLogo") + .arg("-NoProfile") + .arg("-NonInteractive") + .arg("-ExecutionPolicy") + .arg("Bypass") + .arg("-Command") + .arg(command); + shell + } + + #[cfg(not(windows))] + { + let mut shell = Command::new("/bin/bash"); + shell.arg("-c").arg(command); + shell + } +} + +pub fn spawn_shell_command(command: &str, cwd: &Path) -> io::Result { + let mut shell = shell_command(command); + shell + .current_dir(cwd) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + + #[cfg(unix)] + { + use std::os::unix::process::CommandExt; + shell.as_std_mut().process_group(0); + } + + let mut child = shell.spawn()?; + let Some(pid) = child.id() else { + let _ = child.start_kill(); + return Err(io::Error::other( + "spawned command did not expose a process id", + )); + }; + let stdout = child.stdout.take(); + let stderr = child.stderr.take(); + + Ok(SpawnedProcess { + child, + stdout, + stderr, + tree: ProcessTreeGuard::new(pid), + }) +} + +#[derive(Debug)] +struct BoundedBytes { + bytes: Vec, + max_bytes: usize, + truncated: bool, +} + +impl BoundedBytes { + fn new(max_bytes: usize) -> Self { + Self { + bytes: Vec::with_capacity(max_bytes.min(64 * 1024)), + max_bytes, + truncated: false, + } + } + + fn push(&mut self, chunk: &[u8]) { + if self.max_bytes == 0 { + self.truncated |= !chunk.is_empty(); + return; + } + let remaining = self.max_bytes.saturating_sub(self.bytes.len()); + if chunk.len() <= remaining { + self.bytes.extend_from_slice(chunk); + return; + } + self.bytes.extend_from_slice(&chunk[..remaining]); + self.truncated = true; + } + + fn into_text(self) -> (String, bool) { + ( + String::from_utf8_lossy(&self.bytes).into_owned(), + self.truncated, + ) + } +} + +async fn capture_reader(mut reader: R, max_bytes: usize) -> io::Result<(String, bool)> +where + R: AsyncRead + Unpin, +{ + let mut output = BoundedBytes::new(max_bytes); + let mut buffer = vec![0_u8; READ_CHUNK_BYTES]; + loop { + let read = reader.read(&mut buffer).await?; + if read == 0 { + break; + } + output.push(&buffer[..read]); + } + Ok(output.into_text()) +} + +pub async fn run_shell_command( + command: &str, + cwd: &Path, + timeout_ms: u64, + max_capture_bytes: usize, +) -> ProcessRunResult { + let started = Instant::now(); + let mut process = match spawn_shell_command(command, cwd) { + Ok(process) => process, + Err(error) => { + return ProcessRunResult { + stdout: String::new(), + stderr: format!("Failed to execute: {error}"), + success: false, + exit_code: None, + elapsed_ms: started.elapsed().as_millis() as u64, + timed_out: false, + stdout_truncated: false, + stderr_truncated: false, + }; + } + }; + + let stdout_task = process + .take_stdout() + .map(|stdout| tokio::spawn(capture_reader(stdout, max_capture_bytes))); + let stderr_task = process + .take_stderr() + .map(|stderr| tokio::spawn(capture_reader(stderr, max_capture_bytes))); + + let mut timed_out = false; + let status = match timeout(Duration::from_millis(timeout_ms), process.wait()).await { + Ok(Ok(status)) => Some(status), + Ok(Err(error)) => { + process.terminate_tree(); + let _ = process.wait().await; + let mut stderr = format!("Failed while waiting for command: {error}"); + if let Some(task) = stderr_task + && let Ok(Ok((captured, _))) = task.await + && !captured.is_empty() + { + stderr.push('\n'); + stderr.push_str(&captured); + } + let stdout = if let Some(task) = stdout_task { + task.await + .ok() + .and_then(Result::ok) + .map(|entry| entry.0) + .unwrap_or_default() + } else { + String::new() + }; + return ProcessRunResult { + stdout, + stderr, + success: false, + exit_code: None, + elapsed_ms: started.elapsed().as_millis() as u64, + timed_out: false, + stdout_truncated: false, + stderr_truncated: false, + }; + } + Err(_) => { + timed_out = true; + process.terminate_tree(); + process.wait().await.ok() + } + }; + process.disarm(); + + let (stdout, stdout_truncated) = match stdout_task { + Some(task) => task + .await + .ok() + .and_then(Result::ok) + .unwrap_or_else(|| (String::new(), false)), + None => (String::new(), false), + }; + let (mut stderr, stderr_truncated) = match stderr_task { + Some(task) => task + .await + .ok() + .and_then(Result::ok) + .unwrap_or_else(|| (String::new(), false)), + None => (String::new(), false), + }; + + if timed_out { + if !stderr.is_empty() && !stderr.ends_with('\n') { + stderr.push('\n'); + } + stderr.push_str(&format!("Command timed out after {timeout_ms} ms")); + } + + let exit_code = status.as_ref().and_then(std::process::ExitStatus::code); + let success = !timed_out + && status + .as_ref() + .is_some_and(std::process::ExitStatus::success); + + ProcessRunResult { + stdout, + stderr, + success, + exit_code, + elapsed_ms: started.elapsed().as_millis() as u64, + timed_out, + stdout_truncated, + stderr_truncated, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + use uuid::Uuid; + + fn workspace(name: &str) -> PathBuf { + let path = std::env::temp_dir().join(format!("catdesk-process-{name}-{}", Uuid::new_v4())); + std::fs::create_dir_all(&path).expect("create test workspace"); + path + } + + #[tokio::test] + async fn run_shell_command_captures_output_and_exit_status() { + let root = workspace("success"); + let command = if cfg!(windows) { + "Write-Output 'hello'" + } else { + "printf 'hello\\n'" + }; + let result = run_shell_command(command, &root, 5_000, 1024).await; + assert!(result.success, "stderr: {}", result.stderr); + assert_eq!(result.exit_code, Some(0)); + assert_eq!(result.stdout.trim(), "hello"); + assert!(!result.timed_out); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn large_stdout_and_stderr_are_drained_without_deadlock_and_bounded() { + let root = workspace("bounded-output"); + let command = if cfg!(windows) { + "[Console]::Out.Write(('x' * 200000)); [Console]::Error.Write(('y' * 200000))" + } else { + "python3 -c \"import sys; sys.stdout.write('x'*200000); sys.stderr.write('y'*200000)\"" + }; + let result = run_shell_command(command, &root, 5_000, 4_096).await; + assert!( + result.success, + "large-output command failed: {}", + result.stderr + ); + assert!(result.stdout.len() <= 4_096); + assert!(result.stderr.len() <= 4_096); + assert!(result.stdout_truncated); + assert!(result.stderr_truncated); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn timed_out_command_cannot_continue_after_return() { + let root = workspace("timeout"); + let sentinel = root.join("sentinel.txt"); + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 700; Set-Content -Path sentinel.txt -Value survived" + } else { + "sleep 0.7; printf survived > sentinel.txt" + }; + let result = run_shell_command(command, &root, 100, 1024).await; + assert!(result.timed_out); + tokio::time::sleep(Duration::from_millis(900)).await; + assert!( + !sentinel.exists(), + "timed-out process survived and wrote sentinel" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn timeout_terminates_descendant_process_tree() { + let root = workspace("descendant-timeout"); + let sentinel = root.join("descendant.txt"); + let command = if cfg!(windows) { + "Start-Process powershell.exe -ArgumentList '-NoProfile','-Command','Start-Sleep -Milliseconds 800; Set-Content -Path descendant.txt -Value survived' -WorkingDirectory .; Start-Sleep -Seconds 5" + } else { + "(sleep 0.8; printf survived > descendant.txt) & sleep 5" + }; + let result = run_shell_command(command, &root, 150, 1024).await; + assert!(result.timed_out); + tokio::time::sleep(Duration::from_millis(1_000)).await; + assert!( + !sentinel.exists(), + "timed-out root shell left a descendant process alive" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn successful_root_exit_cannot_leave_detached_descendant_alive() { + let root = workspace("detached-success"); + let sentinel = root.join("detached.txt"); + let command = if cfg!(windows) { + "Start-Process powershell.exe -ArgumentList '-NoProfile','-Command','Start-Sleep -Milliseconds 800; Set-Content -Path detached.txt -Value survived' -WorkingDirectory .; Write-Output root-done" + } else { + "(sleep 0.8; printf survived > detached.txt) & printf 'root-done\\n'" + }; + let result = run_shell_command(command, &root, 5_000, 1024).await; + assert!(result.success, "root command failed: {}", result.stderr); + assert!(result.stdout.contains("root-done")); + tokio::time::sleep(Duration::from_millis(1_000)).await; + assert!( + !sentinel.exists(), + "successful root shell detached a descendant outside CatDesk ownership" + ); + let _ = std::fs::remove_dir_all(root); + } + + #[tokio::test] + async fn dropping_run_future_terminates_the_process() { + let root = workspace("drop"); + let sentinel = root.join("sentinel.txt"); + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 700; Set-Content -Path sentinel.txt -Value survived" + } else { + "sleep 0.7; printf survived > sentinel.txt" + }; + let root_for_task = root.clone(); + let task = + tokio::spawn( + async move { run_shell_command(command, &root_for_task, 5_000, 1024).await }, + ); + tokio::time::sleep(Duration::from_millis(100)).await; + task.abort(); + let _ = task.await; + tokio::time::sleep(Duration::from_millis(900)).await; + assert!( + !sentinel.exists(), + "dropped command future left the process alive" + ); + let _ = std::fs::remove_dir_all(root); + } +} diff --git a/src/server.rs b/src/server.rs index e639c72..2ccc812 100644 --- a/src/server.rs +++ b/src/server.rs @@ -11,6 +11,7 @@ use std::collections::HashMap; use std::sync::Arc; use tokio::sync::{Mutex, mpsc::UnboundedSender}; +use crate::command_jobs::CommandJobManager; use crate::devtools::DevtoolsBridge; use crate::mcp::{self, JsonRpcRequest, WIDGET_PAYLOAD_META_KEY}; use crate::state::{ @@ -26,6 +27,7 @@ const STATELESS_FLOW_LABEL: &str = "stateless"; struct ServerState { app: SharedState, devtools: Option>>, + command_jobs: CommandJobManager, ui_events: UnboundedSender, } @@ -33,12 +35,14 @@ struct ServerState { pub fn router( app_state: SharedState, devtools: Option>>, + command_jobs: CommandJobManager, mcp_path: String, ui_events: UnboundedSender, ) -> Router { let state = ServerState { app: app_state, devtools, + command_jobs, ui_events, }; Router::new() @@ -952,6 +956,122 @@ mod tests { ); } + #[tokio::test] + async fn background_command_survives_separate_stateless_http_requests() { + let workspace_root = unique_temp_path("catdesk-post-mcp-command-job"); + let config_root = unique_temp_path("catdesk-post-mcp-command-job-config"); + let config_path = config_root.join("config.toml"); + std::fs::create_dir_all(&workspace_root).expect("create workspace"); + std::fs::create_dir_all(&config_root).expect("create config dir"); + + let app = AppState::new_for_test( + 8787, + workspace_root.to_string_lossy().into_owned(), + config_path.clone(), + ) + .expect("create app state"); + let app_state = Arc::new(Mutex::new(app)); + let (ui_tx, _ui_rx) = unbounded_channel(); + let command_jobs = CommandJobManager::new(); + let server_state = ServerState { + app: app_state, + devtools: None, + command_jobs: command_jobs.clone(), + ui_events: ui_tx, + }; + let command = if cfg!(windows) { + "Start-Sleep -Milliseconds 250; Write-Output http-job-done" + } else { + "sleep 0.25; printf 'http-job-done\\n'" + }; + + let start_response = post_mcp( + State(server_state.clone()), + tool_call_body( + "start_command", + json!({ "command": command, "timeout": 5_000 }), + ), + ) + .await; + assert_eq!(start_response.status(), StatusCode::OK); + let start_body = to_bytes(start_response.into_body(), usize::MAX) + .await + .expect("read start response"); + let start_payload: Value = + serde_json::from_slice(&start_body).expect("parse start response"); + let job_id = start_payload + .get("result") + .and_then(|result| result.get("structuredContent")) + .and_then(|structured| structured.get("jobId")) + .and_then(Value::as_str) + .expect("start response job id") + .to_string(); + + let mut cursor = 0u64; + let mut seen = String::new(); + let mut terminal = None; + for _ in 0..20 { + let poll_response = post_mcp( + State(server_state.clone()), + tool_call_body( + "poll_command", + json!({ "job_id": job_id, "after": cursor, "wait_ms": 250 }), + ), + ) + .await; + assert_eq!(poll_response.status(), StatusCode::OK); + let poll_body = to_bytes(poll_response.into_body(), usize::MAX) + .await + .expect("read poll response"); + let poll_payload: Value = + serde_json::from_slice(&poll_body).expect("parse poll response"); + let structured = poll_payload + .get("result") + .and_then(|result| result.get("structuredContent")) + .expect("poll structured content"); + if let Some(events) = structured.get("events").and_then(Value::as_array) { + for event in events { + if let Some(text) = event.get("text").and_then(Value::as_str) { + seen.push_str(text); + } + } + } + cursor = structured + .get("nextCursor") + .and_then(Value::as_u64) + .unwrap_or(cursor); + let state = structured.get("state").and_then(Value::as_str); + let has_more = structured + .get("hasMoreOutput") + .and_then(Value::as_bool) + .unwrap_or(false); + if state == Some("succeeded") && !has_more { + terminal = Some(poll_payload); + break; + } + } + + let terminal = terminal.expect("background job did not finish across HTTP requests"); + let structured = terminal + .get("result") + .and_then(|result| result.get("structuredContent")) + .expect("terminal structured content"); + assert_eq!( + structured.get("state").and_then(Value::as_str), + Some("succeeded") + ); + assert_eq!( + structured.get("commandSuccess").and_then(Value::as_bool), + Some(true) + ); + assert!(seen.contains("http-job-done")); + + command_jobs.cancel_all().await; + let _ = std::fs::remove_file(config_path); + let _ = std::fs::remove_dir_all(workspace_root); + let _ = std::fs::remove_dir_all(config_root); + } + #[tokio::test] async fn post_mcp_accumulates_usage_from_widget_payload_meta() { let workspace_root = unique_temp_path("catdesk-post-mcp-workspace"); @@ -972,6 +1092,7 @@ mod tests { let server_state = ServerState { app: app_state.clone(), devtools: None, + command_jobs: CommandJobManager::new(), ui_events: ui_tx, }; @@ -1118,6 +1239,7 @@ async fn post_mcp(State(s): State, body_bytes: Bytes) -> Response, pub session_usage_totals: UsageTotals, + pub command_jobs: CommandJobManager, config_path: PathBuf, pub server_handle: Option>, pub ngrok_task: Option>, @@ -944,6 +946,7 @@ impl AppState { request_count: 0, usage_by_model: config.usage_by_model, session_usage_totals: UsageTotals::default(), + command_jobs: CommandJobManager::new(), config_path, server_handle: None, ngrok_task: None, diff --git a/src/widget/catdesk_dashboard.html b/src/widget/catdesk_dashboard.html index 8adf3cd..39e09de 100644 --- a/src/widget/catdesk_dashboard.html +++ b/src/widget/catdesk_dashboard.html @@ -3130,7 +3130,12 @@ return validPayload(next); } - if (toolName === "run_command") { + if ( + toolName === "run_command" || + toolName === "start_command" || + toolName === "poll_command" || + toolName === "cancel_command" + ) { var command = readString(data.command); var commandElapsedMs = readNonNegativeInteger(data.elapsedMs); var commandOutput = readString(data.output); @@ -6549,7 +6554,11 @@ view.panelMode === "tool_call" && view.toolName === "catdesk_instruction"; var isRunCommand = - view.panelMode === "tool_call" && view.toolName === "run_command"; + view.panelMode === "tool_call" && + (view.toolName === "run_command" || + view.toolName === "start_command" || + view.toolName === "poll_command" || + view.toolName === "cancel_command"); var isSearchText = view.panelMode === "tool_call" && view.toolName === "search"; var isReadFile = @@ -6807,6 +6816,9 @@ !nextView.hasChanges && nextView.toolName !== "catdesk_instruction" && nextView.toolName !== "run_command" && + nextView.toolName !== "start_command" && + nextView.toolName !== "poll_command" && + nextView.toolName !== "cancel_command" && nextView.toolName !== "search" && nextView.toolName !== "read" && nextView.toolName !== "list_files" && From 20539218ec811df808aeb9e413d55b10dece1933 Mon Sep 17 00:00:00 2001 From: Navneet Chaudhary Date: Mon, 17 Aug 2026 17:47:40 +0530 Subject: [PATCH 2/3] fix: address async command review feedback --- Cargo.toml | 2 +- src/command_jobs.rs | 123 ++++++++--- src/mcp.rs | 197 +++++++++++++---- src/process_runner.rs | 356 ++++++++++++++++++++++-------- src/widget/catdesk_dashboard.html | 5 +- 5 files changed, 528 insertions(+), 155 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 955539a..1333b47 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,7 @@ ignore = "0.4" libc = "0.2" [target.'cfg(windows)'.dependencies] -windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Security", "Win32_System_JobObjects", "Win32_System_Threading"] } +windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Security", "Win32_System_Diagnostics_ToolHelp", "Win32_System_JobObjects", "Win32_System_Threading"] } [profile.release] lto = true diff --git a/src/command_jobs.rs b/src/command_jobs.rs index 232200b..17df575 100644 --- a/src/command_jobs.rs +++ b/src/command_jobs.rs @@ -20,8 +20,10 @@ const MAX_RETAINED_JOBS: usize = 64; const TERMINAL_JOB_TTL: StdDuration = StdDuration::from_secs(60 * 60); const IDEMPOTENCY_WINDOW: StdDuration = StdDuration::from_secs(30); const MAX_OUTPUT_BYTES_PER_JOB: usize = 4 * 1024 * 1024; +const MAX_TERMINAL_OUTPUT_BYTES: usize = 32 * 1024 * 1024; const MAX_POLL_OUTPUT_BYTES: usize = 128 * 1024; const READ_CHUNK_BYTES: usize = 8 * 1024; +const CLEANUP_INTERVAL: StdDuration = StdDuration::from_secs(1); #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] #[serde(rename_all = "snake_case")] @@ -212,14 +214,6 @@ impl CommandJob { timeout_ms: self.timeout_ms, } } - - async fn terminal_age(&self) -> Option { - self.runtime - .lock() - .await - .finished_at - .map(|finished| finished.elapsed()) - } } #[derive(Default)] @@ -228,6 +222,7 @@ struct ManagerState { // Retry dedupe is intentionally short-lived. JSON-RPC request IDs are only // correlation IDs and may be reused later by a stateless client. request_jobs: HashMap, + last_cleanup: Option, } #[derive(Clone, Default)] @@ -349,15 +344,18 @@ impl CommandJobManager { let job = self.get_job(job_id).await?; let wait_ms = wait_ms.min(MAX_POLL_WAIT_MS); - // Register the notification future before the first snapshot so output - // arriving between the check and wait cannot be missed. + // `Notify::notified()` does not register with `notify_waiters()` until + // the future is polled or explicitly enabled. Pin and enable it before + // taking the snapshot so a change in the check/wait gap is retained. let notified = job.changed.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); let snapshot = job.snapshot(after).await; if snapshot.state.is_terminal() || !snapshot.events.is_empty() || wait_ms == 0 { return Ok(snapshot); } - let _ = timeout(Duration::from_millis(wait_ms), notified).await; + let _ = timeout(Duration::from_millis(wait_ms), &mut notified).await; Ok(job.snapshot(after).await) } @@ -376,9 +374,11 @@ impl CommandJobManager { // misleading Running state. Wait until terminal or the bounded deadline. let deadline = Instant::now() + StdDuration::from_secs(5); loop { - // Register before the snapshot so a terminal transition cannot land - // in the check/wait gap. + // Pin and explicitly enable the waiter before checking state so a + // terminal transition cannot be lost in the check/wait gap. let notified = job.changed.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); let snapshot = job.snapshot(0).await; if snapshot.state.is_terminal() { return Ok(snapshot); @@ -387,7 +387,7 @@ impl CommandJobManager { if remaining.is_zero() { return Ok(snapshot); } - if timeout(remaining, notified).await.is_err() { + if timeout(remaining, &mut notified).await.is_err() { return Ok(job.snapshot(0).await); } } @@ -439,6 +439,17 @@ impl CommandJobManager { } pub async fn cleanup(&self) { + { + let mut manager = self.inner.write().await; + if manager + .last_cleanup + .is_some_and(|last| last.elapsed() < CLEANUP_INTERVAL) + { + return; + } + manager.last_cleanup = Some(Instant::now()); + } + let jobs = { let manager = self.inner.read().await; manager @@ -451,24 +462,49 @@ impl CommandJobManager { let mut expired = Vec::new(); let mut terminal = Vec::new(); for (id, job) in jobs { - if let Some(age) = job.terminal_age().await { - if age >= TERMINAL_JOB_TTL { - expired.push(id); - } else { - terminal.push((id, age)); - } + let runtime = job.runtime.lock().await; + let Some(finished_at) = runtime.finished_at else { + continue; + }; + let age = finished_at.elapsed(); + if age >= TERMINAL_JOB_TTL { + expired.push(id); + } else { + terminal.push((id, age, runtime.retained_output_bytes)); } } - terminal.sort_by_key(|(_, age)| std::cmp::Reverse(*age)); + // Oldest terminal jobs are the first eviction candidates both for the + // retained-job count and for the global decoded-output memory budget. + terminal.sort_by_key(|(_, age, _)| std::cmp::Reverse(*age)); let retained_count = { let manager = self.inner.read().await; manager.jobs.len().saturating_sub(expired.len()) }; if retained_count > MAX_RETAINED_JOBS { let overflow = retained_count - MAX_RETAINED_JOBS; - expired.extend(terminal.into_iter().take(overflow).map(|(id, _)| id)); + expired.extend(terminal.iter().take(overflow).map(|(id, _, _)| id.clone())); } + + let already_expired = expired.iter().cloned().collect::>(); + let mut terminal_output_bytes = terminal + .iter() + .filter(|(id, _, _)| !already_expired.contains(id)) + .map(|(_, _, bytes)| *bytes) + .sum::(); + if terminal_output_bytes > MAX_TERMINAL_OUTPUT_BYTES { + for (id, _, bytes) in &terminal { + if terminal_output_bytes <= MAX_TERMINAL_OUTPUT_BYTES { + break; + } + if already_expired.contains(id) || expired.contains(id) { + continue; + } + expired.push(id.clone()); + terminal_output_bytes = terminal_output_bytes.saturating_sub(*bytes); + } + } + let mut manager = self.inner.write().await; for id in &expired { manager.jobs.remove(id); @@ -558,7 +594,7 @@ where } async fn run_job(job: Arc, mut cancel_rx: watch::Receiver) { - let mut process = match process_runner::spawn_shell_command(&job.command, &job.cwd) { + let mut process = match process_runner::spawn_shell_command(&job.command, &job.cwd).await { Ok(process) => process, Err(error) => { job.append_output("stderr", format!("Failed to execute: {error}\n").as_bytes()) @@ -589,7 +625,7 @@ async fn run_job(job: Arc, mut cancel_rx: watch::Receiver) { let (state, exit_code) = match completion { Completion::Exited(Ok(status)) => { - process.disarm(); + process.disarm().await; if status.success() { (CommandJobState::Succeeded, status.code()) } else { @@ -597,7 +633,7 @@ async fn run_job(job: Arc, mut cancel_rx: watch::Receiver) { } } Completion::Exited(Err(error)) => { - process.terminate_tree(); + process.terminate_tree().await; let _ = process.wait().await; job.append_output( "stderr", @@ -607,7 +643,7 @@ async fn run_job(job: Arc, mut cancel_rx: watch::Receiver) { (CommandJobState::Failed, None) } Completion::Cancelled => { - process.terminate_tree(); + process.terminate_tree().await; let status = process.wait().await.ok(); ( CommandJobState::Cancelled, @@ -615,7 +651,7 @@ async fn run_job(job: Arc, mut cancel_rx: watch::Receiver) { ) } Completion::TimedOut => { - process.terminate_tree(); + process.terminate_tree().await; let status = process.wait().await.ok(); job.append_output( "stderr", @@ -1011,6 +1047,7 @@ mod tests { .get_mut("expired-request-key") .expect("request key exists before cleanup"); entry.1 = Instant::now() - IDEMPOTENCY_WINDOW - StdDuration::from_secs(1); + state.last_cleanup = None; assert!(state.jobs.contains_key(&started.snapshot.job_id)); } @@ -1265,4 +1302,36 @@ mod tests { manager.cancel_all().await; let _ = std::fs::remove_dir_all(root); } + + #[tokio::test] + async fn cleanup_enforces_global_terminal_output_budget() { + let root = workspace("global-output-budget"); + let manager = CommandJobManager::new(); + + for index in 0..9u64 { + let (job, _cancel_rx) = CommandJob::new("synthetic".into(), root.clone(), 5_000); + { + let mut runtime = job.runtime.lock().await; + runtime.state = CommandJobState::Succeeded; + runtime.finished_at = Some(Instant::now() - StdDuration::from_millis(index)); + runtime.retained_output_bytes = MAX_OUTPUT_BYTES_PER_JOB; + } + manager.inner.write().await.jobs.insert(job.id.clone(), job); + } + + manager.inner.write().await.last_cleanup = None; + manager.cleanup().await; + + let jobs = { + let state = manager.inner.read().await; + state.jobs.values().cloned().collect::>() + }; + let mut retained_bytes = 0usize; + for job in jobs { + retained_bytes = + retained_bytes.saturating_add(job.runtime.lock().await.retained_output_bytes); + } + assert!(retained_bytes <= MAX_TERMINAL_OUTPUT_BYTES); + let _ = std::fs::remove_dir_all(root); + } } diff --git a/src/mcp.rs b/src/mcp.rs index 17825fb..99998ac 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -10,7 +10,8 @@ use tokio::sync::Mutex; use crate::command; use crate::command_jobs::{ - CommandJobManager, CommandJobSnapshot, CommandJobState, MAX_JOB_TIMEOUT_MS, MAX_POLL_WAIT_MS, + CommandJobManager, CommandJobSnapshot, CommandJobState, DEFAULT_JOB_TIMEOUT_MS, + MAX_JOB_TIMEOUT_MS, MAX_POLL_WAIT_MS, }; use crate::devtools::DevtoolsBridge; use crate::mascot; @@ -551,7 +552,15 @@ async fn handle_tools_list( "properties": { "command": { "type": "string", "description": "The shell command to execute" }, "cwd": { "type": "string", "description": "Working directory relative to workspace root or absolute path within it" }, - "timeout": { "type": "integer", "minimum": 1, "maximum": 120000, "description": "Timeout in milliseconds for short commands. Maximum 120000; use start_command for long-running work." } + "timeout": { + "type": "integer", + "minimum": 1, + "maximum": command::MAX_TIMEOUT_MS, + "description": format!( + "Timeout in milliseconds for short commands. Maximum {}; use start_command for long-running work.", + command::MAX_TIMEOUT_MS + ) + } }, "required": ["command"] }, @@ -566,7 +575,16 @@ async fn handle_tools_list( "properties": { "command": { "type": "string", "description": "The shell command to start" }, "cwd": { "type": "string", "description": "Working directory relative to workspace root or absolute path within it" }, - "timeout": { "type": "integer", "minimum": 1, "maximum": MAX_JOB_TIMEOUT_MS, "description": "Maximum command runtime in milliseconds. Defaults to 30 minutes; maximum is 24 hours." } + "timeout": { + "type": "integer", + "minimum": 1, + "maximum": MAX_JOB_TIMEOUT_MS, + "description": format!( + "Maximum command runtime in milliseconds. Defaults to {} ms; maximum is {} ms.", + DEFAULT_JOB_TIMEOUT_MS, + MAX_JOB_TIMEOUT_MS + ) + } }, "required": ["command"] }, @@ -906,6 +924,23 @@ async fn forward_to_devtools( } } +fn format_command_output_events<'a, I>(events: I) -> String +where + I: IntoIterator, +{ + let mut output = String::new(); + for (stream, text) in events { + if stream == "stderr" { + output.push_str("[stderr] "); + } + output.push_str(text); + if !text.ends_with('\n') { + output.push('\n'); + } + } + output +} + fn command_job_output_text(snapshot: &CommandJobSnapshot) -> String { if snapshot.events.is_empty() { return match snapshot.state { @@ -913,16 +948,12 @@ fn command_job_output_text(snapshot: &CommandJobSnapshot) -> String { _ => "(no new output)".to_string(), }; } - let mut output = String::new(); - for event in &snapshot.events { - if event.stream == "stderr" { - output.push_str("[stderr] "); - } - output.push_str(&event.text); - if !event.text.ends_with('\n') { - output.push('\n'); - } - } + let mut output = format_command_output_events( + snapshot + .events + .iter() + .map(|event| (event.stream, event.text.as_str())), + ); if snapshot.has_more_output { output.push_str("[more buffered output available; poll again with nextCursor]\n"); } @@ -2149,15 +2180,15 @@ fn widget_changed_files(widget_context: Option<&AutoWidgetContext>) -> (Vec, + show_detail_mode: ShowDetailMode, ) -> Map { let mut payload = Map::new(); let token_stats_layout = current_token_stats_layout(); - let show_detail_mode = current_show_detail_mode(); payload.insert("schema".to_string(), json!("catdesk.review.v1")); payload.insert("panelMode".to_string(), json!(panel_mode)); payload.insert("title".to_string(), json!(title)); @@ -2176,6 +2207,21 @@ fn base_widget_payload( payload } +fn base_widget_payload( + panel_mode: &str, + title: &str, + state: &str, + tool_name: Option<&str>, +) -> Map { + base_widget_payload_with_show_detail_mode( + panel_mode, + title, + state, + tool_name, + current_show_detail_mode(), + ) +} + fn current_token_stats_layout() -> TokenStatsLayout { load_app_config() .map(|config| config.token_stats_layout) @@ -2385,26 +2431,24 @@ fn build_command_job_widget_payload(result: &Value, tool_name: &str) -> Option ("Command Timed Out", "failed"), _ => ("Command Job", "waiting"), }; - let mut output = String::new(); - if let Some(events) = structured.get("events").and_then(Value::as_array) { - for event in events { - let stream = event - .get("stream") - .and_then(Value::as_str) - .unwrap_or("stdout"); - let text = event - .get("text") - .and_then(Value::as_str) - .unwrap_or_default(); - if stream == "stderr" { - output.push_str("[stderr] "); - } - output.push_str(text); - if !text.ends_with('\n') { - output.push('\n'); - } - } - } + let mut output = structured + .get("events") + .and_then(Value::as_array) + .map(|events| { + format_command_output_events(events.iter().map(|event| { + ( + event + .get("stream") + .and_then(Value::as_str) + .unwrap_or("stdout"), + event + .get("text") + .and_then(Value::as_str) + .unwrap_or_default(), + ) + })) + }) + .unwrap_or_default(); if output.is_empty() { output = format!( "job {} ยท {}", @@ -2584,12 +2628,13 @@ fn build_auto_widget_payload( } } -fn enrich_tool_result( +fn enrich_tool_result_with_show_detail_mode( req: &JsonRpcRequest, mut result: Value, widget_context: Option<&AutoWidgetContext>, + show_detail_mode: ShowDetailMode, ) -> Value { - if current_show_detail_mode() == ShowDetailMode::Disable { + if show_detail_mode == ShowDetailMode::Disable { return result; } @@ -2611,7 +2656,14 @@ fn enrich_tool_result( let widget_payload = if has_widget_payload { None } else { - Some(build_auto_widget_payload(req, &result, widget_context)) + let mut payload = build_auto_widget_payload(req, &result, widget_context); + if let Some(payload_obj) = payload.as_object_mut() { + payload_obj.insert( + "showDetailMode".to_string(), + json!(show_detail_mode.as_str()), + ); + } + Some(payload) }; if let Some(result_obj) = result.as_object_mut() { let meta_value = result_obj @@ -2626,6 +2678,19 @@ fn enrich_tool_result( result } +fn enrich_tool_result( + req: &JsonRpcRequest, + result: Value, + widget_context: Option<&AutoWidgetContext>, +) -> Value { + enrich_tool_result_with_show_detail_mode( + req, + result, + widget_context, + current_show_detail_mode(), + ) +} + fn collect_watch_targets(req: &JsonRpcRequest, workspace_root: &str) -> Vec { let tool_name = tool_name_from_request(req); let arguments = tool_arguments(req); @@ -5395,4 +5460,60 @@ hello world" ); assert!(widget_payload.get("widgetMascot").is_some()); } + + #[test] + fn show_detail_modes_are_injectable_for_widget_enrichment() { + let req = tool_call_request("unknown_tool", json!({})); + let raw = json!({ + "content": [{ "type": "text", "text": "hello" }], + "structuredContent": { "toolName": "unknown_tool" } + }); + + let disabled = enrich_tool_result_with_show_detail_mode( + &req, + raw.clone(), + None, + ShowDetailMode::Disable, + ); + assert_eq!( + disabled, raw, + "Disable must leave the tool result untouched" + ); + + for (mode, expected) in [ + (ShowDetailMode::Expanded, "expanded"), + (ShowDetailMode::Collapsed, "collapsed"), + ] { + let result = enrich_tool_result_with_show_detail_mode(&req, raw.clone(), None, mode); + let payload = result + .get("_meta") + .and_then(|meta| meta.get(WIDGET_PAYLOAD_META_KEY)) + .expect("missing injected widget payload"); + assert_eq!( + payload.get("showDetailMode").and_then(Value::as_str), + Some(expected) + ); + } + } + + #[test] + fn base_widget_payload_serializes_all_show_detail_modes() { + for (mode, expected) in [ + (ShowDetailMode::Expanded, "expanded"), + (ShowDetailMode::Collapsed, "collapsed"), + (ShowDetailMode::Disable, "disable"), + ] { + let payload = base_widget_payload_with_show_detail_mode( + "tool_call", + "Test", + "done", + Some("read"), + mode, + ); + assert_eq!( + payload.get("showDetailMode").and_then(Value::as_str), + Some(expected) + ); + } + } } diff --git a/src/process_runner.rs b/src/process_runner.rs index ea1314e..95e916b 100644 --- a/src/process_runner.rs +++ b/src/process_runner.rs @@ -47,9 +47,9 @@ impl SpawnedProcess { } /// Terminate the root process and all descendants owned by this command. - pub fn terminate_tree(&mut self) { - self.tree.terminate(); - // `taskkill /T` / process-group termination should already include the + pub async fn terminate_tree(&mut self) { + self.tree.terminate().await; + // Job-object / process-group termination should already include the // root, but keep Tokio's direct kill as a best-effort fallback. let _ = self.child.start_kill(); } @@ -57,15 +57,15 @@ impl SpawnedProcess { /// Finalize ownership after the root process exits. Any descendants still /// alive at that point are terminated so a command cannot silently detach /// work that outlives its CatDesk job. - pub fn disarm(&mut self) { - self.tree.disarm(); + pub async fn disarm(&mut self) { + self.tree.disarm().await; } } impl Drop for SpawnedProcess { fn drop(&mut self) { if self.tree.is_armed() { - self.tree.terminate(); + self.tree.terminate_blocking(); let _ = self.child.start_kill(); } } @@ -80,12 +80,17 @@ struct ProcessTreeGuard { } impl ProcessTreeGuard { + #[cfg(not(windows))] fn new(pid: u32) -> Self { + Self { pid, armed: true } + } + + #[cfg(windows)] + fn with_windows_job(pid: u32, job_handle: usize) -> Self { Self { pid, armed: true, - #[cfg(windows)] - job_handle: create_windows_job_for_process(pid), + job_handle: Some(job_handle), } } @@ -93,7 +98,7 @@ impl ProcessTreeGuard { self.armed } - fn disarm(&mut self) { + async fn disarm(&mut self) { if !self.armed { return; } @@ -102,8 +107,7 @@ impl ProcessTreeGuard { if self.job_handle.is_some() { close_windows_job(&mut self.job_handle); } else { - // Best effort when Job Object assignment was unavailable. - terminate_process_tree(self.pid); + terminate_process_tree_async(self.pid).await; } } #[cfg(not(windows))] @@ -111,14 +115,29 @@ impl ProcessTreeGuard { self.armed = false; } - fn terminate(&mut self) { + async fn terminate(&mut self) { if !self.armed { return; } #[cfg(windows)] { if !terminate_windows_job(&mut self.job_handle) { - terminate_process_tree(self.pid); + terminate_process_tree_async(self.pid).await; + } + } + #[cfg(not(windows))] + terminate_process_tree(self.pid); + self.armed = false; + } + + fn terminate_blocking(&mut self) { + if !self.armed { + return; + } + #[cfg(windows)] + { + if !terminate_windows_job(&mut self.job_handle) { + terminate_process_tree_blocking(self.pid); } } #[cfg(not(windows))] @@ -129,12 +148,12 @@ impl ProcessTreeGuard { impl Drop for ProcessTreeGuard { fn drop(&mut self) { - self.terminate(); + self.terminate_blocking(); } } #[cfg(windows)] -fn create_windows_job_for_process(pid: u32) -> Option { +fn create_windows_job_for_process(pid: u32) -> io::Result { use std::ffi::c_void; use std::mem::{size_of, zeroed}; use windows_sys::Win32::Foundation::CloseHandle; @@ -150,7 +169,7 @@ fn create_windows_job_for_process(pid: u32) -> Option { unsafe { let job = CreateJobObjectW(std::ptr::null(), std::ptr::null()); if job.is_null() { - return None; + return Err(io::Error::last_os_error()); } let mut info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = zeroed(); @@ -162,23 +181,80 @@ fn create_windows_job_for_process(pid: u32) -> Option { size_of::() as u32, ) == 0 { + let error = io::Error::last_os_error(); CloseHandle(job); - return None; + return Err(error); } let process = OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, 0, pid); if process.is_null() { + let error = io::Error::last_os_error(); CloseHandle(job); - return None; + return Err(error); } let assigned = AssignProcessToJobObject(job, process) != 0; + let assign_error = if assigned { + None + } else { + Some(io::Error::last_os_error()) + }; CloseHandle(process); - if !assigned { + if let Some(error) = assign_error { CloseHandle(job); - return None; + return Err(error); } - Some(job as usize) + Ok(job as usize) + } +} + +#[cfg(windows)] +fn resume_windows_process(pid: u32) -> io::Result<()> { + use std::mem::{size_of, zeroed}; + use windows_sys::Win32::Foundation::{CloseHandle, INVALID_HANDLE_VALUE}; + use windows_sys::Win32::System::Diagnostics::ToolHelp::{ + CreateToolhelp32Snapshot, TH32CS_SNAPTHREAD, THREADENTRY32, Thread32First, Thread32Next, + }; + use windows_sys::Win32::System::Threading::{OpenThread, ResumeThread, THREAD_SUSPEND_RESUME}; + + unsafe { + let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0); + if snapshot == INVALID_HANDLE_VALUE { + return Err(io::Error::last_os_error()); + } + + let mut entry: THREADENTRY32 = zeroed(); + entry.dwSize = size_of::() as u32; + let mut found = Thread32First(snapshot, &mut entry) != 0; + while found { + if entry.th32OwnerProcessID == pid { + let thread = OpenThread(THREAD_SUSPEND_RESUME, 0, entry.th32ThreadID); + if thread.is_null() { + let error = io::Error::last_os_error(); + CloseHandle(snapshot); + return Err(error); + } + let previous_suspend_count = ResumeThread(thread); + let resume_error = if previous_suspend_count == u32::MAX { + Some(io::Error::last_os_error()) + } else { + None + }; + CloseHandle(thread); + CloseHandle(snapshot); + return match resume_error { + Some(error) => Err(error), + None => Ok(()), + }; + } + found = Thread32Next(snapshot, &mut entry) != 0; + } + + CloseHandle(snapshot); + Err(io::Error::new( + io::ErrorKind::NotFound, + "suspended process did not expose a resumable thread", + )) } } @@ -208,10 +284,11 @@ fn terminate_windows_job(job_handle: &mut Option) -> bool { } #[cfg(windows)] -fn terminate_process_tree(pid: u32) { +fn terminate_process_tree_blocking(pid: u32) { // `/T` includes descendants and `/F` makes cancellation deterministic. // Use the executable directly rather than a shell command so the PID never - // passes through shell parsing. + // passes through shell parsing. This synchronous path is reserved for Drop, + // where Rust cannot await cleanup. let _ = std::process::Command::new("taskkill") .args(["/PID", &pid.to_string(), "/T", "/F"]) .stdin(Stdio::null()) @@ -220,6 +297,11 @@ fn terminate_process_tree(pid: u32) { .status(); } +#[cfg(windows)] +async fn terminate_process_tree_async(pid: u32) { + let _ = tokio::task::spawn_blocking(move || terminate_process_tree_blocking(pid)).await; +} + #[cfg(unix)] fn terminate_process_tree(pid: u32) { let pgid = match i32::try_from(pid) { @@ -259,7 +341,7 @@ fn shell_command(command: &str) -> Command { } } -pub fn spawn_shell_command(command: &str, cwd: &Path) -> io::Result { +fn spawn_shell_command_blocking(command: &str, cwd: &Path) -> io::Result { let mut shell = shell_command(command); shell .current_dir(cwd) @@ -274,6 +356,13 @@ pub fn spawn_shell_command(command: &str, cwd: &Path) -> io::Result io::Result handle, + Err(error) => { + let _ = child.start_kill(); + return Err(io::Error::new( + error.kind(), + format!("failed to assign suspended command to Windows Job Object: {error}"), + )); + } + }; + if let Err(error) = resume_windows_process(pid) { + let mut job_handle = Some(job_handle); + close_windows_job(&mut job_handle); + let _ = child.start_kill(); + return Err(io::Error::new( + error.kind(), + format!("failed to resume suspended command process: {error}"), + )); + } + ProcessTreeGuard::with_windows_job(pid, job_handle) + }; + + #[cfg(not(windows))] + let tree = ProcessTreeGuard::new(pid); + let stdout = child.stdout.take(); let stderr = child.stderr.take(); @@ -288,10 +405,18 @@ pub fn spawn_shell_command(command: &str, cwd: &Path) -> io::Result io::Result { + let command = command.to_owned(); + let cwd = cwd.to_path_buf(); + tokio::task::spawn_blocking(move || spawn_shell_command_blocking(&command, &cwd)) + .await + .map_err(|error| io::Error::other(format!("command spawn task failed: {error}")))? +} + #[derive(Debug)] struct BoundedBytes { bytes: Vec, @@ -330,20 +455,59 @@ impl BoundedBytes { } } -async fn capture_reader(mut reader: R, max_bytes: usize) -> io::Result<(String, bool)> +#[derive(Debug, Default)] +struct CapturedOutput { + text: String, + truncated: bool, + read_error: Option, +} + +async fn capture_reader(mut reader: R, max_bytes: usize) -> CapturedOutput where R: AsyncRead + Unpin, { let mut output = BoundedBytes::new(max_bytes); let mut buffer = vec![0_u8; READ_CHUNK_BYTES]; + let mut read_error = None; loop { - let read = reader.read(&mut buffer).await?; - if read == 0 { - break; + match reader.read(&mut buffer).await { + Ok(0) => break, + Ok(read) => output.push(&buffer[..read]), + Err(error) => { + read_error = Some(error.to_string()); + break; + } } - output.push(&buffer[..read]); } - Ok(output.into_text()) + let (text, truncated) = output.into_text(); + CapturedOutput { + text, + truncated, + read_error, + } +} + +async fn finish_capture( + task: Option>, + stream: &str, +) -> CapturedOutput { + let Some(task) = task else { + return CapturedOutput::default(); + }; + match task.await { + Ok(captured) => captured, + Err(error) => CapturedOutput { + read_error: Some(format!("{stream} capture task failed: {error}")), + ..CapturedOutput::default() + }, + } +} + +fn append_stderr_diagnostic(stderr: &mut String, message: &str) { + if !stderr.is_empty() && !stderr.ends_with('\n') { + stderr.push('\n'); + } + stderr.push_str(message); } pub async fn run_shell_command( @@ -353,7 +517,7 @@ pub async fn run_shell_command( max_capture_bytes: usize, ) -> ProcessRunResult { let started = Instant::now(); - let mut process = match spawn_shell_command(command, cwd) { + let mut process = match spawn_shell_command(command, cwd).await { Ok(process) => process, Err(error) => { return ProcessRunResult { @@ -377,73 +541,55 @@ pub async fn run_shell_command( .map(|stderr| tokio::spawn(capture_reader(stderr, max_capture_bytes))); let mut timed_out = false; + let mut wait_error = None; let status = match timeout(Duration::from_millis(timeout_ms), process.wait()).await { Ok(Ok(status)) => Some(status), Ok(Err(error)) => { - process.terminate_tree(); - let _ = process.wait().await; - let mut stderr = format!("Failed while waiting for command: {error}"); - if let Some(task) = stderr_task - && let Ok(Ok((captured, _))) = task.await - && !captured.is_empty() - { - stderr.push('\n'); - stderr.push_str(&captured); - } - let stdout = if let Some(task) = stdout_task { - task.await - .ok() - .and_then(Result::ok) - .map(|entry| entry.0) - .unwrap_or_default() - } else { - String::new() - }; - return ProcessRunResult { - stdout, - stderr, - success: false, - exit_code: None, - elapsed_ms: started.elapsed().as_millis() as u64, - timed_out: false, - stdout_truncated: false, - stderr_truncated: false, - }; + wait_error = Some(error.to_string()); + process.terminate_tree().await; + process.wait().await.ok() } Err(_) => { timed_out = true; - process.terminate_tree(); + process.terminate_tree().await; process.wait().await.ok() } }; - process.disarm(); - - let (stdout, stdout_truncated) = match stdout_task { - Some(task) => task - .await - .ok() - .and_then(Result::ok) - .unwrap_or_else(|| (String::new(), false)), - None => (String::new(), false), - }; - let (mut stderr, stderr_truncated) = match stderr_task { - Some(task) => task - .await - .ok() - .and_then(Result::ok) - .unwrap_or_else(|| (String::new(), false)), - None => (String::new(), false), - }; + process.disarm().await; + + let stdout_capture = finish_capture(stdout_task, "stdout").await; + let stderr_capture = finish_capture(stderr_task, "stderr").await; + let stdout = stdout_capture.text; + let mut stderr = stderr_capture.text; + if let Some(error) = wait_error.as_deref() { + append_stderr_diagnostic( + &mut stderr, + &format!("Failed while waiting for command: {error}"), + ); + } + if let Some(error) = stdout_capture.read_error.as_deref() { + append_stderr_diagnostic( + &mut stderr, + &format!("CatDesk failed to read stdout: {error}"), + ); + } + if let Some(error) = stderr_capture.read_error.as_deref() { + append_stderr_diagnostic( + &mut stderr, + &format!("CatDesk failed to read stderr: {error}"), + ); + } if timed_out { - if !stderr.is_empty() && !stderr.ends_with('\n') { - stderr.push('\n'); - } - stderr.push_str(&format!("Command timed out after {timeout_ms} ms")); + append_stderr_diagnostic( + &mut stderr, + &format!("Command timed out after {timeout_ms} ms"), + ); } let exit_code = status.as_ref().and_then(std::process::ExitStatus::code); - let success = !timed_out + let success = wait_error.is_none() + && !timed_out && status .as_ref() .is_some_and(std::process::ExitStatus::success); @@ -455,8 +601,8 @@ pub async fn run_shell_command( exit_code, elapsed_ms: started.elapsed().as_millis() as u64, timed_out, - stdout_truncated, - stderr_truncated, + stdout_truncated: stdout_capture.truncated, + stderr_truncated: stderr_capture.truncated, } } @@ -464,14 +610,48 @@ pub async fn run_shell_command( mod tests { use super::*; use std::path::PathBuf; + use std::pin::Pin; + use std::task::{Context, Poll}; + use tokio::io::ReadBuf; use uuid::Uuid; + struct PartialThenError { + emitted: bool, + } + + impl AsyncRead for PartialThenError { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if !self.emitted { + self.emitted = true; + buf.put_slice(b"partial-output"); + Poll::Ready(Ok(())) + } else { + Poll::Ready(Err(io::Error::other("synthetic read failure"))) + } + } + } + fn workspace(name: &str) -> PathBuf { let path = std::env::temp_dir().join(format!("catdesk-process-{name}-{}", Uuid::new_v4())); std::fs::create_dir_all(&path).expect("create test workspace"); path } + #[tokio::test] + async fn capture_reader_preserves_partial_output_on_read_error() { + let captured = capture_reader(PartialThenError { emitted: false }, 1024).await; + assert_eq!(captured.text, "partial-output"); + assert!(!captured.truncated); + assert_eq!( + captured.read_error.as_deref(), + Some("synthetic read failure") + ); + } + #[tokio::test] async fn run_shell_command_captures_output_and_exit_status() { let root = workspace("success"); @@ -494,7 +674,7 @@ mod tests { let command = if cfg!(windows) { "[Console]::Out.Write(('x' * 200000)); [Console]::Error.Write(('y' * 200000))" } else { - "python3 -c \"import sys; sys.stdout.write('x'*200000); sys.stderr.write('y'*200000)\"" + "printf '%*s' 200000 ''; printf '%*s' 200000 '' >&2" }; let result = run_shell_command(command, &root, 5_000, 4_096).await; assert!( diff --git a/src/widget/catdesk_dashboard.html b/src/widget/catdesk_dashboard.html index 39e09de..799d445 100644 --- a/src/widget/catdesk_dashboard.html +++ b/src/widget/catdesk_dashboard.html @@ -3139,9 +3139,12 @@ var command = readString(data.command); var commandElapsedMs = readNonNegativeInteger(data.elapsedMs); var commandOutput = readString(data.output); + if (command === null && next.call && next.detail) { + return validPayload(next); + } if (command === null) { return invalidPayload( - "run_command widget payload is missing required fields.", + "command widget payload is missing required fields.", ); } next.command = command; From 017d682ad9ad6bd3af2c1fb5fceab35f717f86f6 Mon Sep 17 00:00:00 2001 From: Navneet Chaudhary Date: Mon, 17 Aug 2026 17:57:49 +0530 Subject: [PATCH 3/3] fix: avoid spawning pre-cancelled jobs --- src/command_jobs.rs | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/src/command_jobs.rs b/src/command_jobs.rs index 17df575..7e2dea4 100644 --- a/src/command_jobs.rs +++ b/src/command_jobs.rs @@ -594,6 +594,12 @@ where } async fn run_job(job: Arc, mut cancel_rx: watch::Receiver) { + let cancelled_before_spawn = *cancel_rx.borrow(); + if cancelled_before_spawn { + job.finish(CommandJobState::Cancelled, None).await; + return; + } + let mut process = match process_runner::spawn_shell_command(&job.command, &job.cwd).await { Ok(process) => process, Err(error) => { @@ -828,6 +834,30 @@ mod tests { let _ = std::fs::remove_dir_all(root); } + #[tokio::test] + async fn pre_cancelled_job_does_not_spawn_command() { + let root = workspace("pre-cancel"); + let sentinel = root.join("sentinel.txt"); + let command = if cfg!(windows) { + "Set-Content sentinel.txt spawned; Start-Sleep -Seconds 2" + } else { + "printf spawned > sentinel.txt; sleep 2" + }; + let (job, cancel_rx) = CommandJob::new(command.to_string(), root.clone(), 10_000); + + let _ = job.cancel_tx.send(true); + run_job(job.clone(), cancel_rx).await; + + let snapshot = job.snapshot(0).await; + assert_eq!(snapshot.state, CommandJobState::Cancelled); + tokio::time::sleep(Duration::from_millis(300)).await; + assert!( + !sentinel.exists(), + "a job cancelled before the runner started still spawned its shell" + ); + let _ = std::fs::remove_dir_all(root); + } + #[tokio::test] async fn cancel_waits_for_terminal_state_despite_output_notifications() { let root = workspace("cancel-terminal");