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            if record.container_state.as_deref()
324                == Some(crate::container_removal::CONTAINER_STATE_REMOVED)
325            {
326                // The container is gone, so there is nothing to ask Docker (#243).
327                return (!process_ids.is_empty()).then_some(Value::Object(process_ids));
328            }
329            let format = "{{.Id}} {{.State.Pid}}";
330            let result = crate::docker_inspect_batch::prefetched_inspect(format, session_name)
331                .unwrap_or_else(|| {
332                    let inspect_args = vec![
333                        "inspect".to_string(),
334                        "-f".to_string(),
335                        format.to_string(),
336                        session_name.to_string(),
337                    ];
338                    let docker = docker_command().to_string_lossy().to_string();
339                    runner.run(&docker, &inspect_args)
340                });
341            if result.success && !result.stdout.trim().is_empty() {
342                let mut parts = result.stdout.split_whitespace();
343                if let Some(container_id) = parts.next() {
344                    insert_if_present(
345                        &mut process_ids,
346                        "containerId",
347                        Some(Value::String(container_id.to_string())),
348                    );
349                }
350                if let Some(pid_value) = parts.next().and_then(parse_pid) {
351                    insert_if_present(&mut process_ids, "containerPid", Some(json!(pid_value)));
352                }
353            }
354        }
355        "ssh" => {
356            insert_if_present(
357                &mut process_ids,
358                "remotePid",
359                record.options.get("remotePid").cloned(),
360            );
361        }
362        _ => {}
363    }
364
365    (!process_ids.is_empty()).then_some(Value::Object(process_ids))
366}
367
368pub fn get_control_command(
369    record: &ExecutionRecord,
370    action: ControlAction,
371) -> Result<ControlCommand, String> {
372    let session_name = record
373        .options
374        .get("sessionName")
375        .and_then(|value| value.as_str())
376        .ok_or_else(|| {
377            "Execution record does not contain an isolation session name.".to_string()
378        })?;
379
380    let isolation_mode = record
381        .options
382        .get("isolationMode")
383        .and_then(|value| value.as_str());
384    if isolation_mode != Some("detached") {
385        return Err("Only detached isolated executions can be stopped or terminated.".to_string());
386    }
387
388    let backend = record
389        .options
390        .get("isolated")
391        .and_then(|value| value.as_str())
392        .unwrap_or("unknown");
393
394    let command = match (action, backend) {
395        (ControlAction::Stop, "screen") => ControlCommand {
396            command: "screen".to_string(),
397            args: vec![
398                "-S".to_string(),
399                session_name.to_string(),
400                "-X".to_string(),
401                "stuff".to_string(),
402                "\u{3}".to_string(),
403            ],
404            method: "CTRL_C".to_string(),
405            message: format!("Sent CTRL+C to detached screen session: {}", session_name),
406        },
407        (ControlAction::Stop, "tmux") => ControlCommand {
408            command: "tmux".to_string(),
409            args: vec![
410                "send-keys".to_string(),
411                "-t".to_string(),
412                session_name.to_string(),
413                "C-c".to_string(),
414            ],
415            method: "CTRL_C".to_string(),
416            message: format!("Sent CTRL+C to detached tmux session: {}", session_name),
417        },
418        (ControlAction::Stop, "docker") => ControlCommand {
419            command: "docker".to_string(),
420            args: vec!["stop".to_string(), session_name.to_string()],
421            method: "DOCKER_STOP".to_string(),
422            message: format!(
423                "Requested graceful stop for detached docker container: {}",
424                session_name
425            ),
426        },
427        (ControlAction::Terminate, "screen") => ControlCommand {
428            command: "screen".to_string(),
429            args: vec![
430                "-S".to_string(),
431                session_name.to_string(),
432                "-X".to_string(),
433                "quit".to_string(),
434            ],
435            method: "SCREEN_QUIT".to_string(),
436            message: format!("Terminated detached screen session: {}", session_name),
437        },
438        (ControlAction::Terminate, "tmux") => ControlCommand {
439            command: "tmux".to_string(),
440            args: vec![
441                "kill-session".to_string(),
442                "-t".to_string(),
443                session_name.to_string(),
444            ],
445            method: "KILL_SESSION".to_string(),
446            message: format!("Terminated detached tmux session: {}", session_name),
447        },
448        (ControlAction::Terminate, "docker") => ControlCommand {
449            command: "docker".to_string(),
450            args: vec!["kill".to_string(), session_name.to_string()],
451            method: "SIGKILL".to_string(),
452            message: format!("Terminated detached docker container: {}", session_name),
453        },
454        (ControlAction::Stop, other) => {
455            return Err(format!(
456                "Stopping detached {} executions is not supported.",
457                other
458            ));
459        }
460        (ControlAction::Terminate, other) => {
461            return Err(format!(
462                "Terminating detached {} executions is not supported.",
463                other
464            ));
465        }
466    };
467
468    Ok(command)
469}
470
471fn append_links_array(lines: &mut Vec<String>, values: &[Value], indent: usize) {
472    let prefix = " ".repeat(indent);
473    if values.is_empty() {
474        lines.push(format!("{}()", prefix));
475        return;
476    }
477
478    lines.push(format!("{}(", prefix));
479    for value in values {
480        match value {
481            Value::Array(nested) => append_links_array(lines, nested, indent + 2),
482            Value::Object(map) => {
483                for (child_key, child_value) in map {
484                    if !child_value.is_null() {
485                        append_links_value(lines, child_key, child_value, indent + 2);
486                    }
487                }
488            }
489            _ => lines.push(format!(
490                "{}{}",
491                " ".repeat(indent + 2),
492                format_value_for_links_notation(value)
493            )),
494        }
495    }
496    lines.push(format!("{})", prefix));
497}
498
499fn append_links_value(lines: &mut Vec<String>, key: &str, value: &Value, indent: usize) {
500    let prefix = " ".repeat(indent);
501    match value {
502        Value::Object(map) => {
503            lines.push(format!("{}{}", prefix, key));
504            for (child_key, child_value) in map {
505                if !child_value.is_null() {
506                    append_links_value(lines, child_key, child_value, indent + 4);
507                }
508            }
509        }
510        Value::Array(values) => {
511            lines.push(format!("{}{}", prefix, key));
512            append_links_array(lines, values, indent + 2);
513        }
514        _ => {
515            lines.push(format!(
516                "{}{} {}",
517                prefix,
518                key,
519                format_value_for_links_notation(value)
520            ));
521        }
522    }
523}
524
525pub fn format_control_result_as_links_notation(
526    action: ControlAction,
527    identifier: &str,
528    record: &ExecutionRecord,
529    method: &str,
530    process_ids: Option<&Value>,
531    message: &str,
532) -> String {
533    let status = match action {
534        ControlAction::Stop => "signal-sent",
535        ControlAction::Terminate => "terminated",
536    };
537    format_control_result_with_status(
538        action,
539        identifier,
540        record,
541        status,
542        method,
543        process_ids,
544        message,
545    )
546}
547
548fn format_control_result_with_status(
549    action: ControlAction,
550    identifier: &str,
551    record: &ExecutionRecord,
552    status: &str,
553    method: &str,
554    process_ids: Option<&Value>,
555    message: &str,
556) -> String {
557    let backend = record
558        .options
559        .get("isolated")
560        .and_then(|value| value.as_str())
561        .unwrap_or("unknown");
562    let session_name = record
563        .options
564        .get("sessionName")
565        .and_then(|value| value.as_str())
566        .unwrap_or("");
567
568    let mut lines = vec![
569        "executionControl".to_string(),
570        format!("  action {}", escape_for_links_notation(action.as_str())),
571        format!("  identifier {}", escape_for_links_notation(identifier)),
572        format!("  uuid {}", escape_for_links_notation(&record.uuid)),
573        format!("  status {}", escape_for_links_notation(status)),
574        format!("  backend {}", escape_for_links_notation(backend)),
575        format!("  sessionName {}", escape_for_links_notation(session_name)),
576        format!("  method {}", escape_for_links_notation(method)),
577    ];
578
579    if let Some(process_ids) = process_ids {
580        append_links_value(&mut lines, "processIds", process_ids, 2);
581    }
582
583    lines.push(format!("  message {}", escape_for_links_notation(message)));
584    lines.join("\n")
585}
586
587pub fn control_execution(
588    store: Option<&ExecutionStore>,
589    identifier: &str,
590    action: ControlAction,
591) -> ExecutionControlResult {
592    control_execution_with_runner(store, identifier, action, &SystemCommandRunner)
593}
594
595fn has_kill_recovery(record: &ExecutionRecord) -> bool {
596    record
597        .options
598        .get("onKillResume")
599        .is_some_and(|value| value.as_u64().unwrap_or(0) > 0)
600}
601
602/// Record that a stop was requested, for sessions with launch-time kill
603/// recovery (issue #176). Returns the previous marker, to restore on failure.
604fn mark_stop_requested(store: &ExecutionStore, record: &mut ExecutionRecord) -> Option<Value> {
605    if !has_kill_recovery(record) {
606        return None;
607    }
608    let previous = record.options.insert(
609        "stopRequestedAt".to_string(),
610        json!(chrono::Utc::now().to_rfc3339()),
611    );
612    // Best effort: without the marker the session may be recovered once more.
613    let _ = store.save(record);
614    previous
615}
616
617/// Whether a failed docker control command hit a container that has already
618/// exited while launch-time recovery is configured: the stop marker alone then
619/// cancels the resume the watcher may still be waiting to make (issue #181).
620fn is_pending_recovery_stop(record: &ExecutionRecord, result: &CommandRunOutput) -> bool {
621    if record.options.get("isolated").and_then(Value::as_str) != Some("docker")
622        || !has_kill_recovery(record)
623    {
624        return false;
625    }
626    let detail = format!(
627        "{} {}",
628        result.stderr,
629        result.error.as_deref().unwrap_or_default()
630    );
631    detail.to_lowercase().contains("is not running")
632}
633
634fn restore_stop_requested(
635    store: &ExecutionStore,
636    record: &mut ExecutionRecord,
637    previous: Option<Value>,
638) {
639    if !has_kill_recovery(record) {
640        return;
641    }
642    match previous {
643        Some(value) => record.options.insert("stopRequestedAt".to_string(), value),
644        None => record.options.remove("stopRequestedAt"),
645    };
646    // Best effort, see mark_stop_requested.
647    let _ = store.save(record);
648}
649
650pub fn control_execution_with_runner<R: CommandRunner>(
651    store: Option<&ExecutionStore>,
652    identifier: &str,
653    action: ControlAction,
654    runner: &R,
655) -> ExecutionControlResult {
656    let Some(store) = store else {
657        return ExecutionControlResult {
658            success: false,
659            output: None,
660            error: Some("Execution tracking is disabled.".to_string()),
661        };
662    };
663
664    let Some(mut record) = store.get(identifier) else {
665        return ExecutionControlResult {
666            success: false,
667            output: None,
668            error: Some(format!(
669                "No execution found with UUID or session name: {}",
670                identifier
671            )),
672        };
673    };
674
675    let control = match get_control_command(&record, action) {
676        Ok(command) => command,
677        Err(error) => {
678            return ExecutionControlResult {
679                success: false,
680                output: None,
681                error: Some(error),
682            }
683        }
684    };
685
686    // A deliberate stop must not be mistaken for a kill that launch-time
687    // recovery should undo (issue #176): `docker stop` escalates to SIGKILL,
688    // which exits 137 just like the OOM killer.
689    let previous_stop_requested_at = mark_stop_requested(store, &mut record);
690    let result = runner.run(&control.command, &control.args);
691    if !result.success && is_pending_recovery_stop(&record, &result) {
692        // The container already exited and its watcher is waiting out
693        // `--on-kill-resume-delay` (issue #181): `docker kill` has nothing to
694        // signal, but the marker it left cancels the pending resume.
695        let session_name = record
696            .options
697            .get("sessionName")
698            .and_then(Value::as_str)
699            .unwrap_or("")
700            .to_string();
701        let output = format_control_result_with_status(
702            action,
703            identifier,
704            &record,
705            "recovery-cancelled",
706            &control.method,
707            None,
708            &format!(
709                "Cancelled the pending recovery of detached docker container: {}",
710                session_name
711            ),
712        );
713        return ExecutionControlResult {
714            success: true,
715            output: Some(output),
716            error: None,
717        };
718    }
719    if !result.success {
720        restore_stop_requested(store, &mut record, previous_stop_requested_at);
721        let backend = record
722            .options
723            .get("isolated")
724            .and_then(|value| value.as_str())
725            .unwrap_or("unknown");
726        let session_name = record
727            .options
728            .get("sessionName")
729            .and_then(|value| value.as_str())
730            .unwrap_or("");
731        let detail = if !result.stderr.is_empty() {
732            result.stderr
733        } else if let Some(error) = result.error {
734            error
735        } else {
736            format!("exit code {}", result.status.unwrap_or(-1))
737        };
738
739        return ExecutionControlResult {
740            success: false,
741            output: None,
742            error: Some(format!(
743                "Failed to {} {} session \"{}\": {}",
744                action.as_str(),
745                backend,
746                session_name,
747                detail
748            )),
749        };
750    }
751
752    let process_ids = collect_process_ids_with_runner(&record, runner);
753    let output = format_control_result_as_links_notation(
754        action,
755        identifier,
756        &record,
757        &control.method,
758        process_ids.as_ref(),
759        &control.message,
760    );
761
762    ExecutionControlResult {
763        success: true,
764        output: Some(output),
765        error: None,
766    }
767}