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 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
592fn 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 let _ = store.save(record);
604 previous
605}
606
607fn 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 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 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 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}