From f8d5dfdce0800f086f0a3b810f149099c2819c5e Mon Sep 17 00:00:00 2001 From: Chronosphere Reliability Date: Sun, 26 Jul 2026 18:56:30 +0000 Subject: [PATCH 1/6] feat: persist subprocess runtime identity --- src/job_runtime.rs | 261 +++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + src/mcp/tools.rs | 51 ++++++--- 3 files changed, 300 insertions(+), 13 deletions(-) create mode 100644 src/job_runtime.rs diff --git a/src/job_runtime.rs b/src/job_runtime.rs new file mode 100644 index 0000000..cd2408e --- /dev/null +++ b/src/job_runtime.rs @@ -0,0 +1,261 @@ +use crate::engagement::{HistoryStore, JobStatus}; +use anyhow::{Context, Result, bail}; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::process::Command; + +pub const JOB_ID_ENV: &str = "CHRONOSPHERE_JOB_ID"; +pub const JOB_TOKEN_ENV: &str = "CHRONOSPHERE_JOB_TOKEN"; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RuntimeIdentity { + pub job_id: String, + pub pid: u32, + pub process_group_id: Option, + pub token: String, + pub started_at: DateTime, +} + +impl RuntimeIdentity { + pub fn new(job_id: String, pid: u32, token: String) -> Self { + Self { + job_id, + pid, + process_group_id: if cfg!(unix) { Some(pid) } else { None }, + token, + started_at: Utc::now(), + } + } +} + +#[derive(Debug, Clone)] +pub struct RunningProcess { + pub identity: RuntimeIdentity, + pub recovered: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum IdentityState { + Alive, + Gone, + Mismatch, + Unsupported, +} + +#[derive(Debug, Default)] +pub struct ReconcileReport { + pub recovered: Vec, + pub stale_jobs: Vec, + pub completed_jobs: Vec, +} + +pub fn runtime_path(jobs_dir: &Path, job_id: &str) -> PathBuf { + jobs_dir.join(format!("{job_id}.runtime.json")) +} + +pub fn persist(jobs_dir: &Path, identity: &RuntimeIdentity) -> Result<()> { + let body = serde_json::to_vec_pretty(identity).context("serialize process identity")?; + crate::security::write_private_atomic(&runtime_path(jobs_dir, &identity.job_id), &body) +} + +pub fn load(jobs_dir: &Path, job_id: &str) -> Result> { + let path = runtime_path(jobs_dir, job_id); + if !path.exists() { + return Ok(None); + } + let body = std::fs::read(&path).with_context(|| format!("read {}", path.display()))?; + let identity = serde_json::from_slice(&body) + .with_context(|| format!("parse runtime identity {}", path.display()))?; + Ok(Some(identity)) +} + +pub fn remove(jobs_dir: &Path, job_id: &str) { + let path = runtime_path(jobs_dir, job_id); + if let Err(err) = std::fs::remove_file(&path) { + if err.kind() != std::io::ErrorKind::NotFound { + tracing::warn!(?err, path = %path.display(), "remove runtime identity failed"); + } + } +} + +pub fn reconcile_history(history: &mut HistoryStore, jobs_dir: &Path) -> ReconcileReport { + let mut report = ReconcileReport::default(); + let running = history + .recent + .iter() + .filter(|record| record.status == JobStatus::Running) + .cloned() + .collect::>(); + + for record in running { + let status_path = jobs_dir.join(format!("{}.status", record.id)); + if let Ok(raw) = std::fs::read_to_string(&status_path) { + if let Ok(code) = raw.trim().parse::() { + let mut updated = record.clone(); + updated.finished_at = Some(Utc::now()); + updated.exit_code = Some(code); + updated.status = if code == 0 { + JobStatus::Completed + } else { + JobStatus::Failed + }; + history.update(&updated); + remove(jobs_dir, &record.id); + report.completed_jobs.push(record.id); + continue; + } + } + + match load(jobs_dir, &record.id) { + Ok(Some(identity)) => match identity_state(&identity) { + IdentityState::Alive => report.recovered.push(identity), + IdentityState::Gone | IdentityState::Mismatch | IdentityState::Unsupported => { + mark_unknown(history, &record); + remove(jobs_dir, &record.id); + report.stale_jobs.push(record.id); + } + }, + Ok(None) if record.tmux_window.is_some() => {} + Ok(None) => { + mark_unknown(history, &record); + report.stale_jobs.push(record.id); + } + Err(err) => { + tracing::warn!(?err, job_id = %record.id, "runtime identity unreadable"); + mark_unknown(history, &record); + remove(jobs_dir, &record.id); + report.stale_jobs.push(record.id); + } + } + } + report +} + +pub fn recover_running_map( + history: &mut HistoryStore, + jobs_dir: &Path, +) -> HashMap { + reconcile_history(history, jobs_dir) + .recovered + .into_iter() + .map(|identity| { + ( + identity.job_id.clone(), + RunningProcess { + identity, + recovered: true, + }, + ) + }) + .collect() +} + +fn mark_unknown(history: &mut HistoryStore, record: &crate::engagement::JobRecord) { + let mut updated = record.clone(); + updated.status = JobStatus::Unknown; + updated.finished_at = Some(Utc::now()); + updated.exit_code = None; + history.update(&updated); +} + +pub fn identity_state(identity: &RuntimeIdentity) -> IdentityState { + #[cfg(target_os = "linux")] + { + linux_identity_state(identity) + } + #[cfg(not(target_os = "linux"))] + { + let _ = identity; + IdentityState::Unsupported + } +} + +#[cfg(target_os = "linux")] +fn linux_identity_state(identity: &RuntimeIdentity) -> IdentityState { + let proc_dir = PathBuf::from(format!("/proc/{}", identity.pid)); + if !proc_dir.exists() { + return IdentityState::Gone; + } + + let stat = match std::fs::read_to_string(proc_dir.join("stat")) { + Ok(stat) => stat, + Err(_) => return IdentityState::Unsupported, + }; + let Some(close) = stat.rfind(')') else { + return IdentityState::Mismatch; + }; + let fields = stat[close + 1..].split_whitespace().collect::>(); + let process_group_id = fields.get(2).and_then(|value| value.parse::().ok()); + if identity.process_group_id.is_some() && process_group_id != identity.process_group_id { + return IdentityState::Mismatch; + } + + let environment = match std::fs::read(proc_dir.join("environ")) { + Ok(environment) => environment, + Err(_) => return IdentityState::Unsupported, + }; + let expected_job = format!("{JOB_ID_ENV}={}", identity.job_id); + let expected_token = format!("{JOB_TOKEN_ENV}={}", identity.token); + let mut has_job = false; + let mut has_token = false; + for item in environment.split(|byte| *byte == 0) { + has_job |= item == expected_job.as_bytes(); + has_token |= item == expected_token.as_bytes(); + } + if has_job && has_token { + IdentityState::Alive + } else { + IdentityState::Mismatch + } +} + +pub fn signal_process_group( + identity: &RuntimeIdentity, + signal: &str, + require_verified_identity: bool, +) -> Result<()> { + #[cfg(unix)] + { + if require_verified_identity && identity_state(identity) != IdentityState::Alive { + bail!( + "refusing to signal job '{}': process identity no longer matches", + identity.job_id + ); + } + let group = identity.process_group_id.unwrap_or(identity.pid); + let status = Command::new("kill") + .arg(format!("-{signal}")) + .arg(format!("-{group}")) + .status() + .with_context(|| format!("send {signal} to process group {group}"))?; + if !status.success() { + bail!("failed to send {signal} to process group {group}"); + } + Ok(()) + } + #[cfg(not(unix))] + { + let _ = (identity, signal, require_verified_identity); + bail!("process-group signalling is unsupported on this platform") + } +} + +pub fn process_group_exists(identity: &RuntimeIdentity) -> bool { + #[cfg(unix)] + { + let group = identity.process_group_id.unwrap_or(identity.pid); + Command::new("kill") + .arg("-0") + .arg(format!("-{group}")) + .status() + .map(|status| status.success()) + .unwrap_or(false) + } + #[cfg(not(unix))] + { + let _ = identity; + false + } +} diff --git a/src/main.rs b/src/main.rs index 6420dbd..f5a40b2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ mod deploy; mod engagement; mod exec; mod input; +mod job_runtime; mod library; mod mcp; mod path_complete; diff --git a/src/mcp/tools.rs b/src/mcp/tools.rs index 74778cf..f35d81b 100644 --- a/src/mcp/tools.rs +++ b/src/mcp/tools.rs @@ -4,6 +4,7 @@ use super::protocol::McpError; use crate::engagement::{CredKind, CredentialProfile, Engagement, JobRecord, JobStatus, Target}; +use crate::job_runtime::{RunningProcess, RuntimeIdentity}; use crate::library::CommandLibrary; use crate::render::{self, RenderContext}; use crate::{builtin, config}; @@ -25,7 +26,7 @@ pub struct State { pub root: PathBuf, pub engagement: Option, pub library: CommandLibrary, - pub running_pids: HashMap, + pub running_jobs: HashMap, } impl State { @@ -61,7 +62,7 @@ impl State { root, engagement, library, - running_pids: HashMap::new(), + running_jobs: HashMap::new(), }) } @@ -668,6 +669,7 @@ async fn tool_run_command(args: Value, state: Arc>) -> Result>) -> Result>) -> Result>) -> Result>) -> Result>) -> Result { - terminate_child(&mut child, pid).await; + terminate_child(&mut child, &identity_for_task).await; (JobStatus::TimedOut, None) } }; let mut state = state_for_task.lock().await; - state.running_pids.remove(&job_id_for_task); + state.running_jobs.remove(&job_id_for_task); + crate::job_runtime::remove(&jobs_dir_for_task, &job_id_for_task); update_job(&mut state, &job_id_for_task, status, exit_code); }); @@ -858,21 +883,21 @@ fn force_update_job(state: &mut State, job_id: &str, status: JobStatus, code: Op } } -async fn terminate_child(child: &mut tokio::process::Child, pid: u32) { +async fn terminate_child(child: &mut tokio::process::Child, identity: &RuntimeIdentity) { #[cfg(unix)] { - let _ = send_process_group_signal(pid, "TERM").await; + let _ = crate::job_runtime::signal_process_group(identity, "TERM", false); if tokio::time::timeout(std::time::Duration::from_secs(2), child.wait()) .await .is_err() { - let _ = send_process_group_signal(pid, "KILL").await; + let _ = crate::job_runtime::signal_process_group(identity, "KILL", false); let _ = child.wait().await; } } #[cfg(not(unix))] { - let _ = pid; + let _ = identity; let _ = child.kill().await; let _ = child.wait().await; } @@ -1045,9 +1070,9 @@ async fn tool_kill_job(args: Value, state: Arc>) -> Result { let pid = { let state = state.lock().await; state - .running_pids + .running_jobs .get(&job_id) - .copied() + .map(|process| process.identity.pid) .ok_or_else(|| anyhow!("job '{}' is not running in this MCP session", job_id))? }; @@ -1074,7 +1099,7 @@ async fn tool_kill_job(args: Value, state: Arc>) -> Result { } let mut state = state.lock().await; - state.running_pids.remove(&job_id); + state.running_jobs.remove(&job_id); force_update_job(&mut state, &job_id, JobStatus::Cancelled, None); Ok(json!({"job_id": job_id, "status": "cancelled"})) } From 8c96e010db1fe8e7e2cc000176457cd8aae3207e Mon Sep 17 00:00:00 2001 From: Chronosphere Reliability Date: Sun, 26 Jul 2026 18:56:35 +0000 Subject: [PATCH 2/6] fix: reconcile stale jobs during startup --- src/engagement/mod.rs | 8 ++++++++ src/mcp/tools.rs | 12 +++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/engagement/mod.rs b/src/engagement/mod.rs index 9006a03..959b4c8 100644 --- a/src/engagement/mod.rs +++ b/src/engagement/mod.rs @@ -166,6 +166,14 @@ impl Engagement { VariableStore::new() }); let mut history = HistoryStore::open(&Self::history_path(&dir))?; + let jobs_dir = Self::jobs_dir(&dir); + let reconciliation = crate::job_runtime::reconcile_history(&mut history, &jobs_dir); + if !reconciliation.stale_jobs.is_empty() { + tracing::warn!( + count = reconciliation.stale_jobs.len(), + "reconciled stale running jobs" + ); + } let secrets = crate::security::store_secrets(&profiles, &aps, &pivots, &variables); if history.redact_values(&secrets) { tracing::warn!("redacted sensitive values from legacy job history"); diff --git a/src/mcp/tools.rs b/src/mcp/tools.rs index f35d81b..bca4a29 100644 --- a/src/mcp/tools.rs +++ b/src/mcp/tools.rs @@ -55,6 +55,16 @@ impl State { } } }; + let mut engagement = engagement; + let running_jobs = engagement + .as_mut() + .map(|engagement| { + crate::job_runtime::recover_running_map( + &mut engagement.history, + &Engagement::jobs_dir(&engagement.dir), + ) + }) + .unwrap_or_default(); let lib_sources = library_sources(&root, engagement.as_ref()); let paths: Vec<&Path> = lib_sources.iter().map(|p| p.as_path()).collect(); let library = CommandLibrary::load(&paths).context("load library")?; @@ -62,7 +72,7 @@ impl State { root, engagement, library, - running_jobs: HashMap::new(), + running_jobs, }) } From 9cd4dff3db9d75fed07a95cbdf21846b26374111 Mon Sep 17 00:00:00 2001 From: Chronosphere Reliability Date: Sun, 26 Jul 2026 18:56:39 +0000 Subject: [PATCH 3/6] fix: support safe cancellation after restart --- src/mcp/tools.rs | 90 +++++++++++++++++++++++++----------------------- 1 file changed, 46 insertions(+), 44 deletions(-) diff --git a/src/mcp/tools.rs b/src/mcp/tools.rs index bca4a29..5c83228 100644 --- a/src/mcp/tools.rs +++ b/src/mcp/tools.rs @@ -913,24 +913,6 @@ async fn terminate_child(child: &mut tokio::process::Child, identity: &RuntimeId } } -#[cfg(unix)] -async fn send_process_group_signal(pid: u32, signal: &str) -> Result<()> { - let status = Command::new("kill") - .arg(format!("-{}", signal)) - .arg(format!("-{}", pid)) - .status() - .await - .with_context(|| format!("send {} to process group {}", signal, pid))?; - if !status.success() { - return Err(anyhow!( - "failed to send {} to process group {}", - signal, - pid - )); - } - Ok(()) -} - #[derive(Deserialize)] struct TailArgs { job_id: String, @@ -1040,9 +1022,37 @@ async fn tool_grep_job(args: Value, state: Arc>) -> Result { Ok(json!({"matches": matches, "count": matches.len()})) } +fn refresh_recovered_jobs(state: &mut State) { + let stale = state + .running_jobs + .iter() + .filter(|(_, process)| { + process.recovered + && crate::job_runtime::identity_state(&process.identity) + != crate::job_runtime::IdentityState::Alive + }) + .map(|(job_id, _)| job_id.clone()) + .collect::>(); + if stale.is_empty() { + return; + } + let jobs_dir = state + .engagement + .as_ref() + .map(|engagement| Engagement::jobs_dir(&engagement.dir)); + for job_id in stale { + state.running_jobs.remove(&job_id); + if let Some(jobs_dir) = &jobs_dir { + crate::job_runtime::remove(jobs_dir, &job_id); + } + force_update_job(state, &job_id, JobStatus::Unknown, None); + } +} + async fn tool_list_jobs(args: Value, state: Arc>) -> Result { let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(20) as usize; - let s = state.lock().await; + let mut s = state.lock().await; + refresh_recovered_jobs(&mut s); let eng = s .engagement .as_ref() @@ -1077,39 +1087,27 @@ async fn tool_kill_job(args: Value, state: Arc>) -> Result { .and_then(Value::as_str) .ok_or_else(|| anyhow!("missing job_id"))? .to_string(); - let pid = { - let state = state.lock().await; + let process = { + let mut state = state.lock().await; + refresh_recovered_jobs(&mut state); state .running_jobs .get(&job_id) - .map(|process| process.identity.pid) - .ok_or_else(|| anyhow!("job '{}' is not running in this MCP session", job_id))? + .cloned() + .ok_or_else(|| anyhow!("job '{}' is not running", job_id))? }; - #[cfg(unix)] - { - send_process_group_signal(pid, "TERM").await?; - tokio::time::sleep(std::time::Duration::from_millis(500)).await; - let alive = Command::new("kill") - .arg("-0") - .arg(format!("-{}", pid)) - .status() - .await - .map(|status| status.success()) - .unwrap_or(false); - if alive { - send_process_group_signal(pid, "KILL").await?; - } - } - #[cfg(not(unix))] - { - return Err(anyhow!( - "job cancellation is not supported on this platform yet" - )); + crate::job_runtime::signal_process_group(&process.identity, "TERM", process.recovered)?; + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + if crate::job_runtime::process_group_exists(&process.identity) { + crate::job_runtime::signal_process_group(&process.identity, "KILL", false)?; } let mut state = state.lock().await; state.running_jobs.remove(&job_id); + if let Some(engagement) = state.engagement.as_ref() { + crate::job_runtime::remove(&Engagement::jobs_dir(&engagement.dir), &job_id); + } force_update_job(&mut state, &job_id, JobStatus::Cancelled, None); Ok(json!({"job_id": job_id, "status": "cancelled"})) } @@ -1284,8 +1282,10 @@ async fn tool_engagement_switch(args: Value, state: Arc>) -> Result .ok_or_else(|| anyhow!("missing name"))? .to_string(); let mut s = state.lock().await; - let eng = Engagement::load_named(&s.root, &name) + let mut eng = Engagement::load_named(&s.root, &name) .with_context(|| format!("load engagement '{}'", name))?; + s.running_jobs = + crate::job_runtime::recover_running_map(&mut eng.history, &Engagement::jobs_dir(&eng.dir)); s.engagement = Some(eng); s.reload_library(); Ok(json!({"ok": true, "engagement": name})) @@ -1310,6 +1310,8 @@ async fn tool_engagement_new(args: Value, state: Arc>) -> Result Date: Sun, 26 Jul 2026 18:56:43 +0000 Subject: [PATCH 4/6] perf: stream bounded job tail and grep output --- src/log_io.rs | 226 +++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + src/mcp/tools.rs | 91 +++++++++++-------- 3 files changed, 283 insertions(+), 35 deletions(-) create mode 100644 src/log_io.rs diff --git a/src/log_io.rs b/src/log_io.rs new file mode 100644 index 0000000..ab04e78 --- /dev/null +++ b/src/log_io.rs @@ -0,0 +1,226 @@ +use anyhow::{Context, Result}; +use serde::Serialize; +use std::collections::VecDeque; +use std::fs::File; +use std::io::{BufRead, BufReader, Read, Seek, SeekFrom}; +use std::path::Path; + +pub const DEFAULT_TAIL_BYTES: u64 = 1024 * 1024; +pub const DEFAULT_GREP_BYTES: u64 = 64 * 1024 * 1024; +pub const MAX_LINE_BYTES: usize = 16 * 1024; + +#[derive(Debug, Serialize)] +pub struct TailOutput { + pub lines: Vec, + pub total_lines: Option, + pub file_bytes: u64, + pub scanned_bytes: u64, + pub truncated: bool, +} + +#[derive(Debug, Serialize)] +pub struct LogMatch { + pub line: usize, + pub text: String, + pub line_truncated: bool, +} + +#[derive(Debug, Serialize)] +pub struct GrepOutput { + pub matches: Vec, + pub scanned_bytes: u64, + pub truncated: bool, +} + +pub fn tail_lines(path: &Path, max_lines: usize, max_bytes: u64) -> Result { + let mut file = File::open(path).with_context(|| format!("open {}", path.display()))?; + let file_bytes = file.metadata()?.len(); + let start = file_bytes.saturating_sub(max_bytes.max(1)); + let total_lines = if start == 0 { + Some(count_lines(path)?) + } else { + None + }; + file.seek(SeekFrom::Start(start))?; + let mut reader = BufReader::new(file); + if start > 0 { + discard_partial_line(&mut reader)?; + } + + let mut lines = VecDeque::with_capacity(max_lines.min(5000)); + let mut scanned_bytes = 0u64; + while let Some(line) = read_bounded_line(&mut reader, MAX_LINE_BYTES, u64::MAX)? { + scanned_bytes = scanned_bytes.saturating_add(line.bytes_consumed); + if lines.len() == max_lines { + lines.pop_front(); + } + lines.push_back(line.text); + } + + Ok(TailOutput { + lines: lines.into_iter().collect(), + total_lines, + file_bytes, + scanned_bytes, + truncated: start > 0, + }) +} + +pub fn grep_lines( + path: &Path, + pattern: &str, + ignore_case: bool, + max_matches: usize, + max_bytes: u64, +) -> Result { + let file = File::open(path).with_context(|| format!("open {}", path.display()))?; + let mut reader = BufReader::new(file); + let needle = if ignore_case { + pattern.to_lowercase() + } else { + pattern.to_string() + }; + let mut scanned_bytes = 0u64; + let mut line_number = 0usize; + let mut matches = Vec::new(); + let mut truncated = false; + + while scanned_bytes < max_bytes { + let remaining = max_bytes - scanned_bytes; + let Some(line) = read_bounded_line(&mut reader, MAX_LINE_BYTES, remaining)? else { + break; + }; + scanned_bytes = scanned_bytes.saturating_add(line.bytes_consumed); + line_number += 1; + let haystack = if ignore_case { + line.text.to_lowercase() + } else { + line.text.clone() + }; + if haystack.contains(&needle) { + matches.push(LogMatch { + line: line_number, + text: line.text, + line_truncated: line.output_truncated, + }); + if matches.len() >= max_matches { + truncated = true; + break; + } + } + if line.scan_limit_reached { + truncated = true; + break; + } + } + + if scanned_bytes >= max_bytes { + truncated = true; + } + Ok(GrepOutput { + matches, + scanned_bytes, + truncated, + }) +} + +struct BoundedLine { + text: String, + bytes_consumed: u64, + output_truncated: bool, + scan_limit_reached: bool, +} + +fn read_bounded_line( + reader: &mut R, + max_output_bytes: usize, + max_scan_bytes: u64, +) -> std::io::Result> { + let mut output = Vec::with_capacity(max_output_bytes.min(4096)); + let mut consumed = 0u64; + let mut output_truncated = false; + let mut saw_any = false; + let mut scan_limit_reached = false; + + loop { + if consumed >= max_scan_bytes { + scan_limit_reached = true; + break; + } + let buffer = reader.fill_buf()?; + if buffer.is_empty() { + break; + } + let buffer_len = buffer.len(); + saw_any = true; + let allowed = usize::try_from((max_scan_bytes - consumed).min(usize::MAX as u64)) + .unwrap_or(usize::MAX) + .min(buffer.len()); + let slice = &buffer[..allowed]; + let newline = slice.iter().position(|byte| *byte == b'\n'); + let take = newline.map_or(slice.len(), |index| index + 1); + let remaining_output = max_output_bytes.saturating_sub(output.len()); + let copy = take.min(remaining_output); + output.extend_from_slice(&slice[..copy]); + output_truncated |= copy < take; + reader.consume(take); + consumed += take as u64; + if newline.is_some() { + break; + } + if take == allowed && allowed < buffer_len { + scan_limit_reached = true; + break; + } + } + + if !saw_any { + return Ok(None); + } + while output + .last() + .is_some_and(|byte| matches!(byte, b'\n' | b'\r')) + { + output.pop(); + } + Ok(Some(BoundedLine { + text: String::from_utf8_lossy(&output).into_owned(), + bytes_consumed: consumed, + output_truncated, + scan_limit_reached, + })) +} + +fn discard_partial_line(reader: &mut R) -> std::io::Result<()> { + loop { + let buffer = reader.fill_buf()?; + if buffer.is_empty() { + return Ok(()); + } + if let Some(index) = buffer.iter().position(|byte| *byte == b'\n') { + reader.consume(index + 1); + return Ok(()); + } + let len = buffer.len(); + reader.consume(len); + } +} + +fn count_lines(path: &Path) -> Result { + let mut file = File::open(path).with_context(|| format!("open {}", path.display()))?; + let mut buffer = [0u8; 64 * 1024]; + let mut count = 0usize; + let mut last = None; + loop { + let read = file.read(&mut buffer)?; + if read == 0 { + break; + } + count += buffer[..read].iter().filter(|byte| **byte == b'\n').count(); + last = buffer.get(read - 1).copied(); + } + if last.is_some_and(|byte| byte != b'\n') { + count += 1; + } + Ok(count) +} diff --git a/src/main.rs b/src/main.rs index f5a40b2..5b159e0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,7 @@ mod exec; mod input; mod job_runtime; mod library; +mod log_io; mod mcp; mod path_complete; mod render; diff --git a/src/mcp/tools.rs b/src/mcp/tools.rs index 5c83228..d526b71 100644 --- a/src/mcp/tools.rs +++ b/src/mcp/tools.rs @@ -16,7 +16,6 @@ use serde_json::{Value, json}; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::Arc; -use tokio::fs; use tokio::process::Command; use tokio::sync::Mutex; @@ -261,7 +260,8 @@ pub fn list_tools() -> Value { "type": "object", "properties": { "job_id": {"type": "string"}, - "lines": {"type": "integer", "minimum": 1, "maximum": 5000, "default": 200} + "lines": {"type": "integer", "minimum": 1, "maximum": 5000, "default": 200}, + "max_bytes": {"type": "integer", "minimum": 4096, "maximum": 67108864, "default": 1048576} }, "required": ["job_id"], "additionalProperties": false @@ -275,7 +275,8 @@ pub fn list_tools() -> Value { "properties": { "job_id": {"type": "string"}, "pattern": {"type": "string"}, - "ignore_case": {"type": "boolean", "default": true} + "ignore_case": {"type": "boolean", "default": true}, + "max_bytes": {"type": "integer", "minimum": 4096, "maximum": 268435456, "default": 67108864} }, "required": ["job_id", "pattern"], "additionalProperties": false @@ -917,13 +918,19 @@ async fn terminate_child(child: &mut tokio::process::Child, identity: &RuntimeId struct TailArgs { job_id: String, lines: Option, + max_bytes: Option, } async fn tool_tail_job(args: Value, state: Arc>) -> Result { let args: TailArgs = serde_json::from_value(args).map_err(|err| anyhow!("{}", err))?; let lines_requested = args.lines.unwrap_or(200).clamp(1, 5000); + let max_bytes = args + .max_bytes + .unwrap_or(crate::log_io::DEFAULT_TAIL_BYTES) + .clamp(4096, 64 * 1024 * 1024); let (log_path, status, exit_code, secrets) = { - let state = state.lock().await; + let mut state = state.lock().await; + refresh_recovered_jobs(&mut state); let engagement = state .engagement .as_ref() @@ -946,20 +953,32 @@ async fn tool_tail_job(args: Value, state: Arc>) -> Result { ), ) }; - let body = match log_path { - Some(path) if path.exists() => fs::read_to_string(&path).await.unwrap_or_default(), - _ => String::new(), + let output = match log_path { + Some(path) if path.exists() => { + crate::log_io::tail_lines(&path, lines_requested, max_bytes)? + } + _ => crate::log_io::TailOutput { + lines: Vec::new(), + total_lines: Some(0), + file_bytes: 0, + scanned_bytes: 0, + truncated: false, + }, }; - let body = crate::security::redact_values(&body, &secrets); - let lines: Vec<&str> = body.lines().collect(); - let start = lines.len().saturating_sub(lines_requested); - let tail: Vec = lines[start..].iter().map(|line| line.to_string()).collect(); + let tail = output + .lines + .iter() + .map(|line| crate::security::redact_values(line, &secrets)) + .collect::>(); Ok(json!({ "job_id": args.job_id, "status": format!("{:?}", status).to_lowercase(), "exit_code": exit_code, "shown_lines": tail.len(), - "total_lines": lines.len(), + "total_lines": output.total_lines, + "file_bytes": output.file_bytes, + "scanned_bytes": output.scanned_bytes, + "truncated": output.truncated, "tail": tail.join(" "), })) @@ -970,18 +989,19 @@ struct GrepArgs { job_id: String, pattern: String, ignore_case: Option, + max_bytes: Option, } async fn tool_grep_job(args: Value, state: Arc>) -> Result { let args: GrepArgs = serde_json::from_value(args).map_err(|err| anyhow!("{}", err))?; let ignore_case = args.ignore_case.unwrap_or(true); - let needle = if ignore_case { - args.pattern.to_lowercase() - } else { - args.pattern.clone() - }; + let max_bytes = args + .max_bytes + .unwrap_or(crate::log_io::DEFAULT_GREP_BYTES) + .clamp(4096, 256 * 1024 * 1024); let (log_path, secrets) = { - let state = state.lock().await; + let mut state = state.lock().await; + refresh_recovered_jobs(&mut state); let engagement = state .engagement .as_ref() @@ -1003,23 +1023,24 @@ async fn tool_grep_job(args: Value, state: Arc>) -> Result { ), ) }; - let body = fs::read_to_string(&log_path).await.unwrap_or_default(); - let body = crate::security::redact_values(&body, &secrets); - let mut matches = Vec::new(); - for (index, line) in body.lines().enumerate() { - let haystack = if ignore_case { - line.to_lowercase() - } else { - line.to_string() - }; - if haystack.contains(&needle) { - matches.push(json!({"line": index + 1, "text": line})); - if matches.len() >= 200 { - break; - } - } - } - Ok(json!({"matches": matches, "count": matches.len()})) + let output = crate::log_io::grep_lines(&log_path, &args.pattern, ignore_case, 200, max_bytes)?; + let matches = output + .matches + .into_iter() + .map(|entry| { + json!({ + "line": entry.line, + "text": crate::security::redact_values(&entry.text, &secrets), + "line_truncated": entry.line_truncated, + }) + }) + .collect::>(); + Ok(json!({ + "matches": matches, + "count": matches.len(), + "scanned_bytes": output.scanned_bytes, + "truncated": output.truncated, + })) } fn refresh_recovered_jobs(state: &mut State) { From 591ef40112f479e32444ffd11efaaa409c3cd918 Mon Sep 17 00:00:00 2001 From: Chronosphere Reliability Date: Sun, 26 Jul 2026 18:58:04 +0000 Subject: [PATCH 5/6] test: cover subprocess recovery and large log streaming --- src/job_runtime.rs | 75 ++++++++++++++++++++++++++++++++++++++++++++++ src/log_io.rs | 58 +++++++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+) diff --git a/src/job_runtime.rs b/src/job_runtime.rs index cd2408e..85841db 100644 --- a/src/job_runtime.rs +++ b/src/job_runtime.rs @@ -259,3 +259,78 @@ pub fn process_group_exists(identity: &RuntimeIdentity) -> bool { false } } + +#[cfg(test)] +mod tests { + use super::*; + + #[cfg(target_os = "linux")] + #[test] + fn verifies_owned_process_and_rejects_pid_reuse() { + use std::os::unix::process::CommandExt; + + let job_id = uuid::Uuid::new_v4().to_string(); + let token = uuid::Uuid::new_v4().to_string(); + let mut command = Command::new("bash"); + command + .arg("-lc") + .arg("sleep 30") + .env(JOB_ID_ENV, &job_id) + .env(JOB_TOKEN_ENV, &token); + command.process_group(0); + let mut child = command.spawn().expect("spawn test child"); + let identity = RuntimeIdentity::new(job_id, child.id(), token); + assert_eq!(identity_state(&identity), IdentityState::Alive); + + let mut wrong = identity.clone(); + wrong.token = "wrong-token".into(); + assert_eq!(identity_state(&wrong), IdentityState::Mismatch); + + signal_process_group(&identity, "TERM", true).expect("terminate test child"); + let _ = child.wait(); + assert_ne!(identity_state(&identity), IdentityState::Alive); + } + + #[test] + fn reconciles_stale_runtime_sidecars_without_signalling() { + let root = std::env::temp_dir().join(format!("chrono-runtime-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let history_path = root.join("jobs.jsonl"); + let jobs_dir = root.join("jobs"); + std::fs::create_dir_all(&jobs_dir).unwrap(); + let mut history = HistoryStore::open(&history_path).unwrap(); + let job_id = uuid::Uuid::new_v4().to_string(); + history + .append(&crate::engagement::JobRecord { + id: job_id.clone(), + command_id: Some("test".into()), + command_title: "test".into(), + resolved: "sleep 30".into(), + started_at: Utc::now(), + finished_at: None, + status: JobStatus::Running, + exit_code: None, + tmux_window: None, + log_path: None, + target: None, + profile: None, + ap: None, + pivot: None, + execution: Some("local".into()), + }) + .unwrap(); + let identity = RuntimeIdentity { + job_id: job_id.clone(), + pid: u32::MAX, + process_group_id: Some(u32::MAX), + token: "stale".into(), + started_at: Utc::now(), + }; + persist(&jobs_dir, &identity).unwrap(); + let report = reconcile_history(&mut history, &jobs_dir); + assert_eq!(report.stale_jobs, vec![job_id.clone()]); + assert_eq!(history.recent[0].status, JobStatus::Unknown); + assert!(!runtime_path(&jobs_dir, &job_id).exists()); + let _ = std::fs::remove_dir_all(root); + } +} diff --git a/src/log_io.rs b/src/log_io.rs index ab04e78..53744ce 100644 --- a/src/log_io.rs +++ b/src/log_io.rs @@ -224,3 +224,61 @@ fn count_lines(path: &Path) -> Result { } Ok(count) } + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + #[test] + fn tails_large_logs_with_bounded_memory_window() { + let dir = std::env::temp_dir().join(format!("chrono-log-tail-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("large.log"); + let mut file = File::create(&path).unwrap(); + for index in 0..100_000usize { + writeln!(file, "line-{index:06}-{}", "x".repeat(32)).unwrap(); + } + let result = tail_lines(&path, 25, 64 * 1024).unwrap(); + assert_eq!(result.lines.len(), 25); + assert!(result.lines.last().unwrap().starts_with("line-099999")); + assert!(result.truncated); + assert!(result.scanned_bytes <= 64 * 1024); + assert_eq!(result.total_lines, None); + let _ = std::fs::remove_dir_all(dir); + } + + #[test] + fn streams_grep_and_caps_matches_and_bytes() { + let dir = std::env::temp_dir().join(format!("chrono-log-grep-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("large.log"); + let mut file = File::create(&path).unwrap(); + for index in 0..20_000usize { + let marker = if index % 100 == 0 { " needle" } else { "" }; + writeln!(file, "row-{index:06}{marker}").unwrap(); + } + let result = grep_lines(&path, "needle", true, 10, 4 * 1024 * 1024).unwrap(); + assert_eq!(result.matches.len(), 10); + assert!(result.truncated); + assert!(result.scanned_bytes <= 4 * 1024 * 1024); + let _ = std::fs::remove_dir_all(dir); + } + + #[test] + fn truncates_single_oversized_lines_without_unbounded_output() { + let dir = std::env::temp_dir().join(format!("chrono-log-line-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("oversized.log"); + std::fs::write( + &path, + format!("needle-{}\n", "x".repeat(MAX_LINE_BYTES * 4)), + ) + .unwrap(); + let result = grep_lines(&path, "needle", false, 5, 8 * 1024 * 1024).unwrap(); + assert_eq!(result.matches.len(), 1); + assert!(result.matches[0].line_truncated); + assert!(result.matches[0].text.len() <= MAX_LINE_BYTES); + let _ = std::fs::remove_dir_all(dir); + } +} From 0980c263dc419656595ed7a4fa65a20a138a9594 Mon Sep 17 00:00:00 2001 From: Chronosphere Reliability Date: Sun, 26 Jul 2026 18:58:09 +0000 Subject: [PATCH 6/6] feat: expand doctor with stale job and orphan cleanup --- src/cli.rs | 54 ++++++++----- src/health.rs | 205 +++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + src/mcp/tools.rs | 46 ++++++----- 4 files changed, 266 insertions(+), 40 deletions(-) create mode 100644 src/health.rs diff --git a/src/cli.rs b/src/cli.rs index ff11682..489e7d4 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -133,6 +133,9 @@ pub enum Command { /// Only print missing tools (script-friendly). #[arg(long)] missing: bool, + /// Reconcile stale jobs and clean orphaned runtime artifacts. + #[arg(long)] + repair: bool, }, /// Extract embedded built-in commands to the user data dir. @@ -566,6 +569,8 @@ engagement_cli!( DoctorCli, "doctor", #[arg(long)] pub missing: bool, + #[arg(long)] + pub repair: bool, ); engagement_cli!( @@ -852,7 +857,10 @@ pub async fn try_early_dispatch() -> Result { let c = DoctorCli::parse_from(&argv); dispatch(cli_from( c.engagement, - Command::Doctor { missing: c.missing }, + Command::Doctor { + missing: c.missing, + repair: c.repair, + }, )) .await? } @@ -1041,30 +1049,38 @@ alias chronosphere='{bin}' } Ok(true) } - Command::Doctor { missing } => { + Command::Doctor { missing, repair } => { let sources = library_sources(root.as_path(), cli.opts.engagement.as_deref())?; let lib = load_library(&sources)?; - let tools = lib.all_tools_referenced(); - let mut found = 0usize; - let mut not_found = Vec::new(); - for tool in &tools { - if which::which(tool).is_ok() { - found += 1; - } else { - not_found.push(tool.clone()); - } - } - not_found.sort(); + let tools = crate::health::check_tools(lib.all_tools_referenced()); if missing { - for t in ¬_found { - println!("{}", t); + for tool in &tools.missing { + println!("{}", tool); } } else { - println!("present: {} / {}", found, tools.len()); - println!("missing:"); - for t in ¬_found { - println!(" - {}", t); + println!("present: {}", tools.present.len()); + println!("missing: {}", tools.missing.len()); + for tool in &tools.missing { + println!(" - {}", tool); + } + } + + match open_engagement(&root, cli.opts.engagement.as_deref()) { + Ok(mut engagement) => { + let health = crate::health::inspect_engagement(&mut engagement, repair)?; + println!("engagement: {}", health.engagement); + println!("running jobs: {}", health.running_jobs.len()); + println!("unknown jobs: {}", health.unknown_jobs.len()); + println!("orphan files: {}", health.orphan_files.len()); + if repair { + println!("removed files: {}", health.removed_files.len()); + println!("archived logs: {}", health.archived_logs.len()); + } + } + Err(err) if cli.opts.engagement.is_none() => { + tracing::debug!(?err, "doctor: no unambiguous engagement selected"); } + Err(err) => return Err(err), } Ok(true) } diff --git a/src/health.rs b/src/health.rs new file mode 100644 index 0000000..1580827 --- /dev/null +++ b/src/health.rs @@ -0,0 +1,205 @@ +use crate::engagement::{Engagement, JobStatus}; +use crate::job_runtime::{self, IdentityState}; +use anyhow::{Context, Result}; +use serde::Serialize; +use std::collections::HashSet; +use std::path::Path; + +#[derive(Debug, Serialize)] +pub struct ToolHealth { + pub present: Vec, + pub missing: Vec, +} + +#[derive(Debug, Default, Serialize)] +pub struct EngagementHealth { + pub engagement: String, + pub running_jobs: Vec, + pub unknown_jobs: Vec, + pub orphan_files: Vec, + pub removed_files: Vec, + pub archived_logs: Vec, +} + +pub fn check_tools(tools: impl IntoIterator) -> ToolHealth { + let mut present = Vec::new(); + let mut missing = Vec::new(); + for tool in tools { + if which::which(&tool).is_ok() { + present.push(tool); + } else { + missing.push(tool); + } + } + present.sort(); + missing.sort(); + ToolHealth { present, missing } +} + +pub fn inspect_engagement(engagement: &mut Engagement, repair: bool) -> Result { + let jobs_dir = Engagement::jobs_dir(&engagement.dir); + std::fs::create_dir_all(&jobs_dir).ok(); + let reconciliation = job_runtime::reconcile_history(&mut engagement.history, &jobs_dir); + let known_jobs = engagement + .history + .recent + .iter() + .map(|record| record.id.clone()) + .collect::>(); + let running_jobs = engagement + .history + .recent + .iter() + .filter(|record| record.status == JobStatus::Running) + .map(|record| record.id.clone()) + .collect::>(); + let mut unknown_jobs = engagement + .history + .recent + .iter() + .filter(|record| record.status == JobStatus::Unknown) + .map(|record| record.id.clone()) + .collect::>(); + unknown_jobs.extend(reconciliation.stale_jobs); + unknown_jobs.sort(); + unknown_jobs.dedup(); + + let mut report = EngagementHealth { + engagement: engagement.meta.name.clone(), + running_jobs, + unknown_jobs, + ..EngagementHealth::default() + }; + + for entry in + std::fs::read_dir(&jobs_dir).with_context(|| format!("read {}", jobs_dir.display()))? + { + let entry = entry?; + let path = entry.path(); + if !path.is_file() { + continue; + } + let Some(name) = path.file_name().and_then(|name| name.to_str()) else { + continue; + }; + let Some((job_id, kind)) = classify_artifact(name) else { + continue; + }; + let record = engagement + .history + .recent + .iter() + .find(|record| record.id == job_id); + let disposable = matches!( + kind, + ArtifactKind::Runtime + | ArtifactKind::Password + | ArtifactKind::RemoteScript + | ArtifactKind::Status + ); + let orphan = !known_jobs.contains(job_id) + || record.is_some_and(|record| { + record.status != JobStatus::Running + && (disposable + || matches!(kind, ArtifactKind::Log) && record.log_path.is_none()) + }) + || matches!(kind, ArtifactKind::Runtime) + && job_runtime::load(&jobs_dir, job_id) + .ok() + .flatten() + .is_some_and(|identity| identity_state_not_alive(&identity)); + if !orphan { + continue; + } + report.orphan_files.push(name.to_string()); + if repair { + if matches!(kind, ArtifactKind::Log) { + archive_log(&jobs_dir, &path, &mut report)?; + } else { + std::fs::remove_file(&path) + .with_context(|| format!("remove orphan {}", path.display()))?; + report.removed_files.push(name.to_string()); + } + } + } + report.orphan_files.sort(); + report.removed_files.sort(); + report.archived_logs.sort(); + Ok(report) +} + +fn identity_state_not_alive(identity: &job_runtime::RuntimeIdentity) -> bool { + !matches!(job_runtime::identity_state(identity), IdentityState::Alive) +} + +#[derive(Debug, Clone, Copy)] +enum ArtifactKind { + Runtime, + Password, + RemoteScript, + Status, + Log, +} + +fn classify_artifact(name: &str) -> Option<(&str, ArtifactKind)> { + for (suffix, kind) in [ + (".runtime.json", ArtifactKind::Runtime), + (".sshpass", ArtifactKind::Password), + (".remote.sh", ArtifactKind::RemoteScript), + (".status", ArtifactKind::Status), + (".log", ArtifactKind::Log), + ] { + if let Some(job_id) = name.strip_suffix(suffix) { + return Some((job_id, kind)); + } + } + None +} + +fn archive_log(jobs_dir: &Path, source: &Path, report: &mut EngagementHealth) -> Result<()> { + let archive = jobs_dir.join("orphaned"); + std::fs::create_dir_all(&archive)?; + let name = source + .file_name() + .ok_or_else(|| anyhow::anyhow!("orphan log has no filename"))?; + let mut destination = archive.join(name); + if destination.exists() { + destination = archive.join(format!( + "{}-{}", + uuid::Uuid::new_v4(), + name.to_string_lossy() + )); + } + std::fs::rename(source, &destination).with_context(|| { + format!( + "archive orphan log {} -> {}", + source.display(), + destination.display() + ) + })?; + report + .archived_logs + .push(destination.to_string_lossy().into_owned()); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn doctor_archives_orphan_logs_and_removes_sensitive_artifacts() { + let root = std::env::temp_dir().join(format!("chrono-health-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&root).unwrap(); + let mut engagement = Engagement::create(&root, "lab").unwrap(); + let jobs = Engagement::jobs_dir(&engagement.dir); + std::fs::write(jobs.join("orphan.log"), "evidence").unwrap(); + std::fs::write(jobs.join("orphan.sshpass"), "secret").unwrap(); + let report = inspect_engagement(&mut engagement, true).unwrap(); + assert_eq!(report.removed_files, vec!["orphan.sshpass"]); + assert_eq!(report.archived_logs.len(), 1); + assert!(!jobs.join("orphan.sshpass").exists()); + assert!(!jobs.join("orphan.log").exists()); + let _ = std::fs::remove_dir_all(root); + } +} diff --git a/src/main.rs b/src/main.rs index 5b159e0..5931145 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,6 +7,7 @@ mod cve; mod deploy; mod engagement; mod exec; +mod health; mod input; mod job_runtime; mod library; diff --git a/src/mcp/tools.rs b/src/mcp/tools.rs index d526b71..4163d57 100644 --- a/src/mcp/tools.rs +++ b/src/mcp/tools.rs @@ -393,10 +393,13 @@ pub fn list_tools() -> Value { }, { "name": "doctor", - "description": "Check which tools referenced by the library are installed (via `which`). Useful for picking commands that will actually work on this host.", + "description": "Check installed tools plus stale jobs and orphaned runtime artifacts for the loaded engagement.", "inputSchema": { "type": "object", - "properties": {"missing_only": {"type": "boolean", "default": false}}, + "properties": { + "missing_only": {"type": "boolean", "default": false}, + "repair": {"type": "boolean", "default": false} + }, "additionalProperties": false } } @@ -1341,29 +1344,30 @@ async fn tool_engagement_new(args: Value, state: Arc>) -> Result>) -> Result { let missing_only = args .get("missing_only") - .and_then(|v| v.as_bool()) + .and_then(Value::as_bool) .unwrap_or(false); - let s = state.lock().await; - let tools = s.library.all_tools_referenced(); - let mut present = Vec::new(); - let mut missing = Vec::new(); - for t in tools { - if which::which(&t).is_ok() { - present.push(t); - } else { - missing.push(t); - } - } - present.sort(); - missing.sort(); + let repair = args.get("repair").and_then(Value::as_bool).unwrap_or(false); + let mut state = state.lock().await; + refresh_recovered_jobs(&mut state); + let tools = crate::health::check_tools(state.library.all_tools_referenced()); + let engagement = state + .engagement + .as_mut() + .map(|engagement| crate::health::inspect_engagement(engagement, repair)) + .transpose()?; if missing_only { - Ok(json!({"missing": missing, "missing_count": missing.len()})) + Ok(json!({ + "missing": tools.missing, + "missing_count": tools.missing.len(), + "engagement": engagement, + })) } else { Ok(json!({ - "present": present, - "missing": missing, - "present_count": present.len(), - "missing_count": missing.len(), + "present": tools.present, + "missing": tools.missing, + "present_count": tools.present.len(), + "missing_count": tools.missing.len(), + "engagement": engagement, })) } }