Skip to main content

start_command/
execution_control.rs

1//! Detached execution control helpers.
2//!
3//! Maps tracked detached execution records back to native isolation backend
4//! controls so callers can stop or terminate a running session by UUID or
5//! session name.
6
7use crate::docker_cleanup::docker_command;
8use crate::execution_store::{ExecutionRecord, ExecutionStore};
9use crate::output_blocks::{escape_for_links_notation, format_value_for_links_notation};
10use serde_json::{json, Map, Value};
11use std::collections::{HashSet, VecDeque};
12use std::process::Command;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum ControlAction {
16    Stop,
17    Terminate,
18}
19
20impl ControlAction {
21    pub fn as_str(self) -> &'static str {
22        match self {
23            ControlAction::Stop => "stop",
24            ControlAction::Terminate => "terminate",
25        }
26    }
27}
28
29#[derive(Debug, Clone, Default, PartialEq, Eq)]
30pub struct CommandRunOutput {
31    pub success: bool,
32    pub stdout: String,
33    pub stderr: String,
34    pub status: Option<i32>,
35    pub error: Option<String>,
36}
37
38pub trait CommandRunner {
39    fn run(&self, command: &str, args: &[String]) -> CommandRunOutput;
40
41    /// Run with an upper bound; the process is killed once `timeout` passes.
42    /// Runners that cannot time out simply run the command.
43    fn run_with_timeout(
44        &self,
45        command: &str,
46        args: &[String],
47        timeout: std::time::Duration,
48    ) -> CommandRunOutput {
49        let _ = timeout;
50        self.run(command, args)
51    }
52}
53
54#[derive(Debug, Default)]
55pub struct SystemCommandRunner;
56
57impl CommandRunner for SystemCommandRunner {
58    fn run(&self, command: &str, args: &[String]) -> CommandRunOutput {
59        match Command::new(command).args(args).output() {
60            Ok(output) => CommandRunOutput {
61                success: output.status.success(),
62                stdout: String::from_utf8_lossy(&output.stdout).to_string(),
63                stderr: String::from_utf8_lossy(&output.stderr).to_string(),
64                status: output.status.code(),
65                error: None,
66            },
67            Err(err) => CommandRunOutput {
68                success: false,
69                stdout: String::new(),
70                stderr: String::new(),
71                status: None,
72                error: Some(err.to_string()),
73            },
74        }
75    }
76
77    fn run_with_timeout(
78        &self,
79        command: &str,
80        args: &[String],
81        timeout: std::time::Duration,
82    ) -> CommandRunOutput {
83        use std::io::Read;
84        use std::process::Stdio;
85        let spawned = Command::new(command)
86            .args(args)
87            .stdin(Stdio::null())
88            .stdout(Stdio::piped())
89            .stderr(Stdio::piped())
90            .spawn();
91        let mut child = match spawned {
92            Ok(child) => child,
93            Err(err) => {
94                return CommandRunOutput {
95                    error: Some(err.to_string()),
96                    ..Default::default()
97                }
98            }
99        };
100        let drain = |pipe: Option<Box<dyn Read + Send>>| {
101            std::thread::spawn(move || {
102                let mut text = String::new();
103                if let Some(mut pipe) = pipe {
104                    let mut bytes = Vec::new();
105                    let _ = pipe.read_to_end(&mut bytes);
106                    text = String::from_utf8_lossy(&bytes).to_string();
107                }
108                text
109            })
110        };
111        let stdout = drain(
112            child
113                .stdout
114                .take()
115                .map(|p| Box::new(p) as Box<dyn Read + Send>),
116        );
117        let stderr = drain(
118            child
119                .stderr
120                .take()
121                .map(|p| Box::new(p) as Box<dyn Read + Send>),
122        );
123        let started = std::time::Instant::now();
124        let status = loop {
125            match child.try_wait() {
126                Ok(Some(status)) => break Ok(status),
127                Ok(None) if started.elapsed() >= timeout => {
128                    let _ = child.kill();
129                    let _ = child.wait();
130                    break Err(format!("timed out after {}s", timeout.as_secs_f64()));
131                }
132                Ok(None) => std::thread::sleep(std::time::Duration::from_millis(50)),
133                Err(err) => break Err(err.to_string()),
134            }
135        };
136        let stdout = stdout.join().unwrap_or_default();
137        let stderr = stderr.join().unwrap_or_default();
138        match status {
139            Ok(status) => CommandRunOutput {
140                success: status.success(),
141                stdout,
142                stderr,
143                status: status.code(),
144                error: None,
145            },
146            Err(error) => CommandRunOutput {
147                stdout,
148                stderr,
149                error: Some(error),
150                ..Default::default()
151            },
152        }
153    }
154}
155
156#[derive(Debug, Clone, PartialEq, Eq)]
157pub struct ControlCommand {
158    pub command: String,
159    pub args: Vec<String>,
160    pub method: String,
161    pub message: String,
162}
163
164pub struct ExecutionControlResult {
165    pub success: bool,
166    pub output: Option<String>,
167    pub error: Option<String>,
168}
169
170fn parse_pid(value: &str) -> Option<u32> {
171    value.trim().parse::<u32>().ok().filter(|pid| *pid > 0)
172}
173
174fn parse_pids(output: &str) -> Vec<u32> {
175    output.split_whitespace().filter_map(parse_pid).collect()
176}
177
178pub fn parse_screen_pid(screen_list_output: &str, session_name: &str) -> Option<u32> {
179    for line in screen_list_output.lines() {
180        let first_column = line.split_whitespace().next().unwrap_or("");
181        let Some((pid, name)) = first_column.split_once('.') else {
182            continue;
183        };
184        if name == session_name {
185            return parse_pid(pid);
186        }
187    }
188    None
189}
190
191pub fn collect_descendant_pids_with_runner<R: CommandRunner>(
192    root_pid: u32,
193    runner: &R,
194) -> Vec<u32> {
195    let mut descendants = Vec::new();
196    let mut seen = HashSet::from([root_pid]);
197    let mut queue = VecDeque::from([root_pid]);
198
199    while let Some(parent_pid) = queue.pop_front() {
200        let args = vec!["-P".to_string(), parent_pid.to_string()];
201        let result = runner.run("pgrep", &args);
202        if !result.success && result.stdout.is_empty() {
203            continue;
204        }
205
206        for child_pid in parse_pids(&result.stdout) {
207            if seen.insert(child_pid) {
208                descendants.push(child_pid);
209                queue.push_back(child_pid);
210            }
211        }
212    }
213
214    descendants
215}
216
217pub fn collect_descendant_pids(root_pid: u32) -> Vec<u32> {
218    collect_descendant_pids_with_runner(root_pid, &SystemCommandRunner)
219}
220
221fn insert_if_present(map: &mut Map<String, Value>, key: &str, value: Option<Value>) {
222    if let Some(value) = value {
223        if !matches!(&value, Value::Array(items) if items.is_empty()) {
224            map.insert(key.to_string(), value);
225        }
226    }
227}
228
229fn option_value_as_string(value: Option<&Value>) -> Option<String> {
230    match value {
231        Some(Value::String(s)) => Some(s.clone()),
232        Some(Value::Number(n)) => Some(n.to_string()),
233        _ => None,
234    }
235}
236
237pub fn collect_process_ids(record: &ExecutionRecord) -> Option<Value> {
238    collect_process_ids_with_runner(record, &SystemCommandRunner)
239}
240
241pub fn collect_process_ids_with_runner<R: CommandRunner>(
242    record: &ExecutionRecord,
243    runner: &R,
244) -> Option<Value> {
245    let mut process_ids = Map::new();
246    insert_if_present(
247        &mut process_ids,
248        "wrapperPid",
249        record.pid.map(|pid| json!(pid)),
250    );
251
252    let session_name = record
253        .options
254        .get("sessionName")
255        .and_then(|value| value.as_str());
256    let isolated = record
257        .options
258        .get("isolated")
259        .and_then(|value| value.as_str());
260
261    let (Some(session_name), Some(isolated)) = (session_name, isolated) else {
262        return (!process_ids.is_empty()).then_some(Value::Object(process_ids));
263    };
264
265    match isolated {
266        "screen" => {
267            let result = runner.run("screen", &["-ls".to_string()]);
268            let output = format!("{}{}", result.stdout, result.stderr);
269            if let Some(screen_pid) = parse_screen_pid(&output, session_name) {
270                insert_if_present(&mut process_ids, "screenPid", Some(json!(screen_pid)));
271                insert_if_present(
272                    &mut process_ids,
273                    "commandPids",
274                    Some(json!(collect_descendant_pids_with_runner(
275                        screen_pid, runner
276                    ))),
277                );
278            }
279        }
280        "tmux" => {
281            let tmux_pid_args = vec![
282                "display-message".to_string(),
283                "-p".to_string(),
284                "-t".to_string(),
285                session_name.to_string(),
286                "#{pid}".to_string(),
287            ];
288            let tmux_pid_result = runner.run("tmux", &tmux_pid_args);
289            insert_if_present(
290                &mut process_ids,
291                "tmuxPid",
292                parse_pid(&tmux_pid_result.stdout).map(|pid| json!(pid)),
293            );
294
295            let pane_args = vec![
296                "list-panes".to_string(),
297                "-t".to_string(),
298                session_name.to_string(),
299                "-F".to_string(),
300                "#{pane_pid}".to_string(),
301            ];
302            let pane_result = runner.run("tmux", &pane_args);
303            let pane_pids = parse_pids(&pane_result.stdout);
304            insert_if_present(&mut process_ids, "panePids", Some(json!(pane_pids)));
305
306            let mut command_pids = Vec::new();
307            let mut seen = HashSet::new();
308            for pane_pid in parse_pids(&pane_result.stdout) {
309                for command_pid in collect_descendant_pids_with_runner(pane_pid, runner) {
310                    if seen.insert(command_pid) {
311                        command_pids.push(command_pid);
312                    }
313                }
314            }
315            insert_if_present(&mut process_ids, "commandPids", Some(json!(command_pids)));
316        }
317        "docker" => {
318            insert_if_present(
319                &mut process_ids,
320                "containerId",
321                option_value_as_string(record.options.get("containerId")).map(Value::String),
322            );
323            let inspect_args = vec![
324                "inspect".to_string(),
325                "-f".to_string(),
326                "{{.Id}} {{.State.Pid}}".to_string(),
327                session_name.to_string(),
328            ];
329            let docker = docker_command().to_string_lossy().to_string();
330            let result = runner.run(&docker, &inspect_args);
331            if result.success && !result.stdout.trim().is_empty() {
332                let mut parts = result.stdout.split_whitespace();
333                if let Some(container_id) = parts.next() {
334                    insert_if_present(
335                        &mut process_ids,
336                        "containerId",
337                        Some(Value::String(container_id.to_string())),
338                    );
339                }
340                if let Some(pid_value) = parts.next().and_then(parse_pid) {
341                    insert_if_present(&mut process_ids, "containerPid", Some(json!(pid_value)));
342                }
343            }
344        }
345        "ssh" => {
346            insert_if_present(
347                &mut process_ids,
348                "remotePid",
349                record.options.get("remotePid").cloned(),
350            );
351        }
352        _ => {}
353    }
354
355    (!process_ids.is_empty()).then_some(Value::Object(process_ids))
356}
357
358pub fn get_control_command(
359    record: &ExecutionRecord,
360    action: ControlAction,
361) -> Result<ControlCommand, String> {
362    let session_name = record
363        .options
364        .get("sessionName")
365        .and_then(|value| value.as_str())
366        .ok_or_else(|| {
367            "Execution record does not contain an isolation session name.".to_string()
368        })?;
369
370    let isolation_mode = record
371        .options
372        .get("isolationMode")
373        .and_then(|value| value.as_str());
374    if isolation_mode != Some("detached") {
375        return Err("Only detached isolated executions can be stopped or terminated.".to_string());
376    }
377
378    let backend = record
379        .options
380        .get("isolated")
381        .and_then(|value| value.as_str())
382        .unwrap_or("unknown");
383
384    let command = match (action, backend) {
385        (ControlAction::Stop, "screen") => ControlCommand {
386            command: "screen".to_string(),
387            args: vec![
388                "-S".to_string(),
389                session_name.to_string(),
390                "-X".to_string(),
391                "stuff".to_string(),
392                "\u{3}".to_string(),
393            ],
394            method: "CTRL_C".to_string(),
395            message: format!("Sent CTRL+C to detached screen session: {}", session_name),
396        },
397        (ControlAction::Stop, "tmux") => ControlCommand {
398            command: "tmux".to_string(),
399            args: vec![
400                "send-keys".to_string(),
401                "-t".to_string(),
402                session_name.to_string(),
403                "C-c".to_string(),
404            ],
405            method: "CTRL_C".to_string(),
406            message: format!("Sent CTRL+C to detached tmux session: {}", session_name),
407        },
408        (ControlAction::Stop, "docker") => ControlCommand {
409            command: "docker".to_string(),
410            args: vec!["stop".to_string(), session_name.to_string()],
411            method: "DOCKER_STOP".to_string(),
412            message: format!(
413                "Requested graceful stop for detached docker container: {}",
414                session_name
415            ),
416        },
417        (ControlAction::Terminate, "screen") => ControlCommand {
418            command: "screen".to_string(),
419            args: vec![
420                "-S".to_string(),
421                session_name.to_string(),
422                "-X".to_string(),
423                "quit".to_string(),
424            ],
425            method: "SCREEN_QUIT".to_string(),
426            message: format!("Terminated detached screen session: {}", session_name),
427        },
428        (ControlAction::Terminate, "tmux") => ControlCommand {
429            command: "tmux".to_string(),
430            args: vec![
431                "kill-session".to_string(),
432                "-t".to_string(),
433                session_name.to_string(),
434            ],
435            method: "KILL_SESSION".to_string(),
436            message: format!("Terminated detached tmux session: {}", session_name),
437        },
438        (ControlAction::Terminate, "docker") => ControlCommand {
439            command: "docker".to_string(),
440            args: vec!["kill".to_string(), session_name.to_string()],
441            method: "SIGKILL".to_string(),
442            message: format!("Terminated detached docker container: {}", session_name),
443        },
444        (ControlAction::Stop, other) => {
445            return Err(format!(
446                "Stopping detached {} executions is not supported.",
447                other
448            ));
449        }
450        (ControlAction::Terminate, other) => {
451            return Err(format!(
452                "Terminating detached {} executions is not supported.",
453                other
454            ));
455        }
456    };
457
458    Ok(command)
459}
460
461fn append_links_array(lines: &mut Vec<String>, values: &[Value], indent: usize) {
462    let prefix = " ".repeat(indent);
463    if values.is_empty() {
464        lines.push(format!("{}()", prefix));
465        return;
466    }
467
468    lines.push(format!("{}(", prefix));
469    for value in values {
470        match value {
471            Value::Array(nested) => append_links_array(lines, nested, indent + 2),
472            Value::Object(map) => {
473                for (child_key, child_value) in map {
474                    if !child_value.is_null() {
475                        append_links_value(lines, child_key, child_value, indent + 2);
476                    }
477                }
478            }
479            _ => lines.push(format!(
480                "{}{}",
481                " ".repeat(indent + 2),
482                format_value_for_links_notation(value)
483            )),
484        }
485    }
486    lines.push(format!("{})", prefix));
487}
488
489fn append_links_value(lines: &mut Vec<String>, key: &str, value: &Value, indent: usize) {
490    let prefix = " ".repeat(indent);
491    match value {
492        Value::Object(map) => {
493            lines.push(format!("{}{}", prefix, key));
494            for (child_key, child_value) in map {
495                if !child_value.is_null() {
496                    append_links_value(lines, child_key, child_value, indent + 4);
497                }
498            }
499        }
500        Value::Array(values) => {
501            lines.push(format!("{}{}", prefix, key));
502            append_links_array(lines, values, indent + 2);
503        }
504        _ => {
505            lines.push(format!(
506                "{}{} {}",
507                prefix,
508                key,
509                format_value_for_links_notation(value)
510            ));
511        }
512    }
513}
514
515pub fn format_control_result_as_links_notation(
516    action: ControlAction,
517    identifier: &str,
518    record: &ExecutionRecord,
519    method: &str,
520    process_ids: Option<&Value>,
521    message: &str,
522) -> String {
523    let status = match action {
524        ControlAction::Stop => "signal-sent",
525        ControlAction::Terminate => "terminated",
526    };
527    format_control_result_with_status(
528        action,
529        identifier,
530        record,
531        status,
532        method,
533        process_ids,
534        message,
535    )
536}
537
538fn format_control_result_with_status(
539    action: ControlAction,
540    identifier: &str,
541    record: &ExecutionRecord,
542    status: &str,
543    method: &str,
544    process_ids: Option<&Value>,
545    message: &str,
546) -> String {
547    let backend = record
548        .options
549        .get("isolated")
550        .and_then(|value| value.as_str())
551        .unwrap_or("unknown");
552    let session_name = record
553        .options
554        .get("sessionName")
555        .and_then(|value| value.as_str())
556        .unwrap_or("");
557
558    let mut lines = vec![
559        "executionControl".to_string(),
560        format!("  action {}", escape_for_links_notation(action.as_str())),
561        format!("  identifier {}", escape_for_links_notation(identifier)),
562        format!("  uuid {}", escape_for_links_notation(&record.uuid)),
563        format!("  status {}", escape_for_links_notation(status)),
564        format!("  backend {}", escape_for_links_notation(backend)),
565        format!("  sessionName {}", escape_for_links_notation(session_name)),
566        format!("  method {}", escape_for_links_notation(method)),
567    ];
568
569    if let Some(process_ids) = process_ids {
570        append_links_value(&mut lines, "processIds", process_ids, 2);
571    }
572
573    lines.push(format!("  message {}", escape_for_links_notation(message)));
574    lines.join("\n")
575}
576
577pub fn control_execution(
578    store: Option<&ExecutionStore>,
579    identifier: &str,
580    action: ControlAction,
581) -> ExecutionControlResult {
582    control_execution_with_runner(store, identifier, action, &SystemCommandRunner)
583}
584
585fn has_kill_recovery(record: &ExecutionRecord) -> bool {
586    record
587        .options
588        .get("onKillResume")
589        .is_some_and(|value| value.as_u64().unwrap_or(0) > 0)
590}
591
592/// Record that a stop was requested, for sessions with launch-time kill
593/// recovery (issue #176). Returns the previous marker, to restore on failure.
594fn mark_stop_requested(store: &ExecutionStore, record: &mut ExecutionRecord) -> Option<Value> {
595    if !has_kill_recovery(record) {
596        return None;
597    }
598    let previous = record.options.insert(
599        "stopRequestedAt".to_string(),
600        json!(chrono::Utc::now().to_rfc3339()),
601    );
602    // Best effort: without the marker the session may be recovered once more.
603    let _ = store.save(record);
604    previous
605}
606
607/// Whether a failed docker control command hit a container that has already
608/// exited while launch-time recovery is configured: the stop marker alone then
609/// cancels the resume the watcher may still be waiting to make (issue #181).
610fn is_pending_recovery_stop(record: &ExecutionRecord, result: &CommandRunOutput) -> bool {
611    if record.options.get("isolated").and_then(Value::as_str) != Some("docker")
612        || !has_kill_recovery(record)
613    {
614        return false;
615    }
616    let detail = format!(
617        "{} {}",
618        result.stderr,
619        result.error.as_deref().unwrap_or_default()
620    );
621    detail.to_lowercase().contains("is not running")
622}
623
624fn restore_stop_requested(
625    store: &ExecutionStore,
626    record: &mut ExecutionRecord,
627    previous: Option<Value>,
628) {
629    if !has_kill_recovery(record) {
630        return;
631    }
632    match previous {
633        Some(value) => record.options.insert("stopRequestedAt".to_string(), value),
634        None => record.options.remove("stopRequestedAt"),
635    };
636    // Best effort, see mark_stop_requested.
637    let _ = store.save(record);
638}
639
640pub fn control_execution_with_runner<R: CommandRunner>(
641    store: Option<&ExecutionStore>,
642    identifier: &str,
643    action: ControlAction,
644    runner: &R,
645) -> ExecutionControlResult {
646    let Some(store) = store else {
647        return ExecutionControlResult {
648            success: false,
649            output: None,
650            error: Some("Execution tracking is disabled.".to_string()),
651        };
652    };
653
654    let Some(mut record) = store.get(identifier) else {
655        return ExecutionControlResult {
656            success: false,
657            output: None,
658            error: Some(format!(
659                "No execution found with UUID or session name: {}",
660                identifier
661            )),
662        };
663    };
664
665    let control = match get_control_command(&record, action) {
666        Ok(command) => command,
667        Err(error) => {
668            return ExecutionControlResult {
669                success: false,
670                output: None,
671                error: Some(error),
672            }
673        }
674    };
675
676    // A deliberate stop must not be mistaken for a kill that launch-time
677    // recovery should undo (issue #176): `docker stop` escalates to SIGKILL,
678    // which exits 137 just like the OOM killer.
679    let previous_stop_requested_at = mark_stop_requested(store, &mut record);
680    let result = runner.run(&control.command, &control.args);
681    if !result.success && is_pending_recovery_stop(&record, &result) {
682        // The container already exited and its watcher is waiting out
683        // `--on-kill-resume-delay` (issue #181): `docker kill` has nothing to
684        // signal, but the marker it left cancels the pending resume.
685        let session_name = record
686            .options
687            .get("sessionName")
688            .and_then(Value::as_str)
689            .unwrap_or("")
690            .to_string();
691        let output = format_control_result_with_status(
692            action,
693            identifier,
694            &record,
695            "recovery-cancelled",
696            &control.method,
697            None,
698            &format!(
699                "Cancelled the pending recovery of detached docker container: {}",
700                session_name
701            ),
702        );
703        return ExecutionControlResult {
704            success: true,
705            output: Some(output),
706            error: None,
707        };
708    }
709    if !result.success {
710        restore_stop_requested(store, &mut record, previous_stop_requested_at);
711        let backend = record
712            .options
713            .get("isolated")
714            .and_then(|value| value.as_str())
715            .unwrap_or("unknown");
716        let session_name = record
717            .options
718            .get("sessionName")
719            .and_then(|value| value.as_str())
720            .unwrap_or("");
721        let detail = if !result.stderr.is_empty() {
722            result.stderr
723        } else if let Some(error) = result.error {
724            error
725        } else {
726            format!("exit code {}", result.status.unwrap_or(-1))
727        };
728
729        return ExecutionControlResult {
730            success: false,
731            output: None,
732            error: Some(format!(
733                "Failed to {} {} session \"{}\": {}",
734                action.as_str(),
735                backend,
736                session_name,
737                detail
738            )),
739        };
740    }
741
742    let process_ids = collect_process_ids_with_runner(&record, runner);
743    let output = format_control_result_as_links_notation(
744        action,
745        identifier,
746        &record,
747        &control.method,
748        process_ids.as_ref(),
749        &control.message,
750    );
751
752    ExecutionControlResult {
753        success: true,
754        output: Some(output),
755        error: None,
756    }
757}