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