1use 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 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 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
602fn 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 let _ = store.save(record);
614 previous
615}
616
617fn 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 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 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 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}