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
42#[derive(Debug, Default)]
43pub struct SystemCommandRunner;
44
45impl CommandRunner for SystemCommandRunner {
46 fn run(&self, command: &str, args: &[String]) -> CommandRunOutput {
47 match Command::new(command).args(args).output() {
48 Ok(output) => CommandRunOutput {
49 success: output.status.success(),
50 stdout: String::from_utf8_lossy(&output.stdout).to_string(),
51 stderr: String::from_utf8_lossy(&output.stderr).to_string(),
52 status: output.status.code(),
53 error: None,
54 },
55 Err(err) => CommandRunOutput {
56 success: false,
57 stdout: String::new(),
58 stderr: String::new(),
59 status: None,
60 error: Some(err.to_string()),
61 },
62 }
63 }
64}
65
66#[derive(Debug, Clone, PartialEq, Eq)]
67pub struct ControlCommand {
68 pub command: String,
69 pub args: Vec<String>,
70 pub method: String,
71 pub message: String,
72}
73
74pub struct ExecutionControlResult {
75 pub success: bool,
76 pub output: Option<String>,
77 pub error: Option<String>,
78}
79
80fn parse_pid(value: &str) -> Option<u32> {
81 value.trim().parse::<u32>().ok().filter(|pid| *pid > 0)
82}
83
84fn parse_pids(output: &str) -> Vec<u32> {
85 output.split_whitespace().filter_map(parse_pid).collect()
86}
87
88pub fn parse_screen_pid(screen_list_output: &str, session_name: &str) -> Option<u32> {
89 for line in screen_list_output.lines() {
90 let first_column = line.split_whitespace().next().unwrap_or("");
91 let Some((pid, name)) = first_column.split_once('.') else {
92 continue;
93 };
94 if name == session_name {
95 return parse_pid(pid);
96 }
97 }
98 None
99}
100
101pub fn collect_descendant_pids_with_runner<R: CommandRunner>(
102 root_pid: u32,
103 runner: &R,
104) -> Vec<u32> {
105 let mut descendants = Vec::new();
106 let mut seen = HashSet::from([root_pid]);
107 let mut queue = VecDeque::from([root_pid]);
108
109 while let Some(parent_pid) = queue.pop_front() {
110 let args = vec!["-P".to_string(), parent_pid.to_string()];
111 let result = runner.run("pgrep", &args);
112 if !result.success && result.stdout.is_empty() {
113 continue;
114 }
115
116 for child_pid in parse_pids(&result.stdout) {
117 if seen.insert(child_pid) {
118 descendants.push(child_pid);
119 queue.push_back(child_pid);
120 }
121 }
122 }
123
124 descendants
125}
126
127pub fn collect_descendant_pids(root_pid: u32) -> Vec<u32> {
128 collect_descendant_pids_with_runner(root_pid, &SystemCommandRunner)
129}
130
131fn insert_if_present(map: &mut Map<String, Value>, key: &str, value: Option<Value>) {
132 if let Some(value) = value {
133 if !matches!(&value, Value::Array(items) if items.is_empty()) {
134 map.insert(key.to_string(), value);
135 }
136 }
137}
138
139fn option_value_as_string(value: Option<&Value>) -> Option<String> {
140 match value {
141 Some(Value::String(s)) => Some(s.clone()),
142 Some(Value::Number(n)) => Some(n.to_string()),
143 _ => None,
144 }
145}
146
147pub fn collect_process_ids(record: &ExecutionRecord) -> Option<Value> {
148 collect_process_ids_with_runner(record, &SystemCommandRunner)
149}
150
151pub fn collect_process_ids_with_runner<R: CommandRunner>(
152 record: &ExecutionRecord,
153 runner: &R,
154) -> Option<Value> {
155 let mut process_ids = Map::new();
156 insert_if_present(
157 &mut process_ids,
158 "wrapperPid",
159 record.pid.map(|pid| json!(pid)),
160 );
161
162 let session_name = record
163 .options
164 .get("sessionName")
165 .and_then(|value| value.as_str());
166 let isolated = record
167 .options
168 .get("isolated")
169 .and_then(|value| value.as_str());
170
171 let (Some(session_name), Some(isolated)) = (session_name, isolated) else {
172 return (!process_ids.is_empty()).then_some(Value::Object(process_ids));
173 };
174
175 match isolated {
176 "screen" => {
177 let result = runner.run("screen", &["-ls".to_string()]);
178 let output = format!("{}{}", result.stdout, result.stderr);
179 if let Some(screen_pid) = parse_screen_pid(&output, session_name) {
180 insert_if_present(&mut process_ids, "screenPid", Some(json!(screen_pid)));
181 insert_if_present(
182 &mut process_ids,
183 "commandPids",
184 Some(json!(collect_descendant_pids_with_runner(
185 screen_pid, runner
186 ))),
187 );
188 }
189 }
190 "tmux" => {
191 let tmux_pid_args = vec![
192 "display-message".to_string(),
193 "-p".to_string(),
194 "-t".to_string(),
195 session_name.to_string(),
196 "#{pid}".to_string(),
197 ];
198 let tmux_pid_result = runner.run("tmux", &tmux_pid_args);
199 insert_if_present(
200 &mut process_ids,
201 "tmuxPid",
202 parse_pid(&tmux_pid_result.stdout).map(|pid| json!(pid)),
203 );
204
205 let pane_args = vec![
206 "list-panes".to_string(),
207 "-t".to_string(),
208 session_name.to_string(),
209 "-F".to_string(),
210 "#{pane_pid}".to_string(),
211 ];
212 let pane_result = runner.run("tmux", &pane_args);
213 let pane_pids = parse_pids(&pane_result.stdout);
214 insert_if_present(&mut process_ids, "panePids", Some(json!(pane_pids)));
215
216 let mut command_pids = Vec::new();
217 let mut seen = HashSet::new();
218 for pane_pid in parse_pids(&pane_result.stdout) {
219 for command_pid in collect_descendant_pids_with_runner(pane_pid, runner) {
220 if seen.insert(command_pid) {
221 command_pids.push(command_pid);
222 }
223 }
224 }
225 insert_if_present(&mut process_ids, "commandPids", Some(json!(command_pids)));
226 }
227 "docker" => {
228 insert_if_present(
229 &mut process_ids,
230 "containerId",
231 option_value_as_string(record.options.get("containerId")).map(Value::String),
232 );
233 let inspect_args = vec![
234 "inspect".to_string(),
235 "-f".to_string(),
236 "{{.Id}} {{.State.Pid}}".to_string(),
237 session_name.to_string(),
238 ];
239 let docker = docker_command().to_string_lossy().to_string();
240 let result = runner.run(&docker, &inspect_args);
241 if result.success && !result.stdout.trim().is_empty() {
242 let mut parts = result.stdout.split_whitespace();
243 if let Some(container_id) = parts.next() {
244 insert_if_present(
245 &mut process_ids,
246 "containerId",
247 Some(Value::String(container_id.to_string())),
248 );
249 }
250 if let Some(pid_value) = parts.next().and_then(parse_pid) {
251 insert_if_present(&mut process_ids, "containerPid", Some(json!(pid_value)));
252 }
253 }
254 }
255 "ssh" => {
256 insert_if_present(
257 &mut process_ids,
258 "remotePid",
259 record.options.get("remotePid").cloned(),
260 );
261 }
262 _ => {}
263 }
264
265 (!process_ids.is_empty()).then_some(Value::Object(process_ids))
266}
267
268pub fn get_control_command(
269 record: &ExecutionRecord,
270 action: ControlAction,
271) -> Result<ControlCommand, String> {
272 let session_name = record
273 .options
274 .get("sessionName")
275 .and_then(|value| value.as_str())
276 .ok_or_else(|| {
277 "Execution record does not contain an isolation session name.".to_string()
278 })?;
279
280 let isolation_mode = record
281 .options
282 .get("isolationMode")
283 .and_then(|value| value.as_str());
284 if isolation_mode != Some("detached") {
285 return Err("Only detached isolated executions can be stopped or terminated.".to_string());
286 }
287
288 let backend = record
289 .options
290 .get("isolated")
291 .and_then(|value| value.as_str())
292 .unwrap_or("unknown");
293
294 let command = match (action, backend) {
295 (ControlAction::Stop, "screen") => ControlCommand {
296 command: "screen".to_string(),
297 args: vec![
298 "-S".to_string(),
299 session_name.to_string(),
300 "-X".to_string(),
301 "stuff".to_string(),
302 "\u{3}".to_string(),
303 ],
304 method: "CTRL_C".to_string(),
305 message: format!("Sent CTRL+C to detached screen session: {}", session_name),
306 },
307 (ControlAction::Stop, "tmux") => ControlCommand {
308 command: "tmux".to_string(),
309 args: vec![
310 "send-keys".to_string(),
311 "-t".to_string(),
312 session_name.to_string(),
313 "C-c".to_string(),
314 ],
315 method: "CTRL_C".to_string(),
316 message: format!("Sent CTRL+C to detached tmux session: {}", session_name),
317 },
318 (ControlAction::Stop, "docker") => ControlCommand {
319 command: "docker".to_string(),
320 args: vec!["stop".to_string(), session_name.to_string()],
321 method: "DOCKER_STOP".to_string(),
322 message: format!(
323 "Requested graceful stop for detached docker container: {}",
324 session_name
325 ),
326 },
327 (ControlAction::Terminate, "screen") => ControlCommand {
328 command: "screen".to_string(),
329 args: vec![
330 "-S".to_string(),
331 session_name.to_string(),
332 "-X".to_string(),
333 "quit".to_string(),
334 ],
335 method: "SCREEN_QUIT".to_string(),
336 message: format!("Terminated detached screen session: {}", session_name),
337 },
338 (ControlAction::Terminate, "tmux") => ControlCommand {
339 command: "tmux".to_string(),
340 args: vec![
341 "kill-session".to_string(),
342 "-t".to_string(),
343 session_name.to_string(),
344 ],
345 method: "KILL_SESSION".to_string(),
346 message: format!("Terminated detached tmux session: {}", session_name),
347 },
348 (ControlAction::Terminate, "docker") => ControlCommand {
349 command: "docker".to_string(),
350 args: vec!["kill".to_string(), session_name.to_string()],
351 method: "SIGKILL".to_string(),
352 message: format!("Terminated detached docker container: {}", session_name),
353 },
354 (ControlAction::Stop, other) => {
355 return Err(format!(
356 "Stopping detached {} executions is not supported.",
357 other
358 ));
359 }
360 (ControlAction::Terminate, other) => {
361 return Err(format!(
362 "Terminating detached {} executions is not supported.",
363 other
364 ));
365 }
366 };
367
368 Ok(command)
369}
370
371fn append_links_array(lines: &mut Vec<String>, values: &[Value], indent: usize) {
372 let prefix = " ".repeat(indent);
373 if values.is_empty() {
374 lines.push(format!("{}()", prefix));
375 return;
376 }
377
378 lines.push(format!("{}(", prefix));
379 for value in values {
380 match value {
381 Value::Array(nested) => append_links_array(lines, nested, indent + 2),
382 Value::Object(map) => {
383 for (child_key, child_value) in map {
384 if !child_value.is_null() {
385 append_links_value(lines, child_key, child_value, indent + 2);
386 }
387 }
388 }
389 _ => lines.push(format!(
390 "{}{}",
391 " ".repeat(indent + 2),
392 format_value_for_links_notation(value)
393 )),
394 }
395 }
396 lines.push(format!("{})", prefix));
397}
398
399fn append_links_value(lines: &mut Vec<String>, key: &str, value: &Value, indent: usize) {
400 let prefix = " ".repeat(indent);
401 match value {
402 Value::Object(map) => {
403 lines.push(format!("{}{}", prefix, key));
404 for (child_key, child_value) in map {
405 if !child_value.is_null() {
406 append_links_value(lines, child_key, child_value, indent + 4);
407 }
408 }
409 }
410 Value::Array(values) => {
411 lines.push(format!("{}{}", prefix, key));
412 append_links_array(lines, values, indent + 2);
413 }
414 _ => {
415 lines.push(format!(
416 "{}{} {}",
417 prefix,
418 key,
419 format_value_for_links_notation(value)
420 ));
421 }
422 }
423}
424
425pub fn format_control_result_as_links_notation(
426 action: ControlAction,
427 identifier: &str,
428 record: &ExecutionRecord,
429 method: &str,
430 process_ids: Option<&Value>,
431 message: &str,
432) -> String {
433 let backend = record
434 .options
435 .get("isolated")
436 .and_then(|value| value.as_str())
437 .unwrap_or("unknown");
438 let session_name = record
439 .options
440 .get("sessionName")
441 .and_then(|value| value.as_str())
442 .unwrap_or("");
443 let status = match action {
444 ControlAction::Stop => "signal-sent",
445 ControlAction::Terminate => "terminated",
446 };
447
448 let mut lines = vec![
449 "executionControl".to_string(),
450 format!(" action {}", escape_for_links_notation(action.as_str())),
451 format!(" identifier {}", escape_for_links_notation(identifier)),
452 format!(" uuid {}", escape_for_links_notation(&record.uuid)),
453 format!(" status {}", escape_for_links_notation(status)),
454 format!(" backend {}", escape_for_links_notation(backend)),
455 format!(" sessionName {}", escape_for_links_notation(session_name)),
456 format!(" method {}", escape_for_links_notation(method)),
457 ];
458
459 if let Some(process_ids) = process_ids {
460 append_links_value(&mut lines, "processIds", process_ids, 2);
461 }
462
463 lines.push(format!(" message {}", escape_for_links_notation(message)));
464 lines.join("\n")
465}
466
467pub fn control_execution(
468 store: Option<&ExecutionStore>,
469 identifier: &str,
470 action: ControlAction,
471) -> ExecutionControlResult {
472 control_execution_with_runner(store, identifier, action, &SystemCommandRunner)
473}
474
475pub fn control_execution_with_runner<R: CommandRunner>(
476 store: Option<&ExecutionStore>,
477 identifier: &str,
478 action: ControlAction,
479 runner: &R,
480) -> ExecutionControlResult {
481 let Some(store) = store else {
482 return ExecutionControlResult {
483 success: false,
484 output: None,
485 error: Some("Execution tracking is disabled.".to_string()),
486 };
487 };
488
489 let Some(record) = store.get(identifier) else {
490 return ExecutionControlResult {
491 success: false,
492 output: None,
493 error: Some(format!(
494 "No execution found with UUID or session name: {}",
495 identifier
496 )),
497 };
498 };
499
500 let control = match get_control_command(&record, action) {
501 Ok(command) => command,
502 Err(error) => {
503 return ExecutionControlResult {
504 success: false,
505 output: None,
506 error: Some(error),
507 }
508 }
509 };
510
511 let result = runner.run(&control.command, &control.args);
512 if !result.success {
513 let backend = record
514 .options
515 .get("isolated")
516 .and_then(|value| value.as_str())
517 .unwrap_or("unknown");
518 let session_name = record
519 .options
520 .get("sessionName")
521 .and_then(|value| value.as_str())
522 .unwrap_or("");
523 let detail = if !result.stderr.is_empty() {
524 result.stderr
525 } else if let Some(error) = result.error {
526 error
527 } else {
528 format!("exit code {}", result.status.unwrap_or(-1))
529 };
530
531 return ExecutionControlResult {
532 success: false,
533 output: None,
534 error: Some(format!(
535 "Failed to {} {} session \"{}\": {}",
536 action.as_str(),
537 backend,
538 session_name,
539 detail
540 )),
541 };
542 }
543
544 let process_ids = collect_process_ids_with_runner(&record, runner);
545 let output = format_control_result_as_links_notation(
546 action,
547 identifier,
548 &record,
549 &control.method,
550 process_ids.as_ref(),
551 &control.message,
552 );
553
554 ExecutionControlResult {
555 success: true,
556 output: Some(output),
557 error: None,
558 }
559}