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![
319 "kill".to_string(),
320 "--signal=SIGINT".to_string(),
321 session_name.to_string(),
322 ],
323 method: "SIGINT".to_string(),
324 message: format!("Sent SIGINT to detached docker container: {}", session_name),
325 },
326 (ControlAction::Terminate, "screen") => ControlCommand {
327 command: "screen".to_string(),
328 args: vec![
329 "-S".to_string(),
330 session_name.to_string(),
331 "-X".to_string(),
332 "quit".to_string(),
333 ],
334 method: "SCREEN_QUIT".to_string(),
335 message: format!("Terminated detached screen session: {}", session_name),
336 },
337 (ControlAction::Terminate, "tmux") => ControlCommand {
338 command: "tmux".to_string(),
339 args: vec![
340 "kill-session".to_string(),
341 "-t".to_string(),
342 session_name.to_string(),
343 ],
344 method: "KILL_SESSION".to_string(),
345 message: format!("Terminated detached tmux session: {}", session_name),
346 },
347 (ControlAction::Terminate, "docker") => ControlCommand {
348 command: "docker".to_string(),
349 args: vec!["kill".to_string(), session_name.to_string()],
350 method: "SIGKILL".to_string(),
351 message: format!("Terminated detached docker container: {}", session_name),
352 },
353 (ControlAction::Stop, other) => {
354 return Err(format!(
355 "Stopping detached {} executions is not supported.",
356 other
357 ));
358 }
359 (ControlAction::Terminate, other) => {
360 return Err(format!(
361 "Terminating detached {} executions is not supported.",
362 other
363 ));
364 }
365 };
366
367 Ok(command)
368}
369
370fn append_links_value(lines: &mut Vec<String>, key: &str, value: &Value, indent: usize) {
371 let prefix = " ".repeat(indent);
372 match value {
373 Value::Object(map) => {
374 lines.push(format!("{}{}", prefix, key));
375 for (child_key, child_value) in map {
376 if !child_value.is_null() {
377 append_links_value(lines, child_key, child_value, indent + 2);
378 }
379 }
380 }
381 _ => {
382 lines.push(format!(
383 "{}{} {}",
384 prefix,
385 key,
386 format_value_for_links_notation(value)
387 ));
388 }
389 }
390}
391
392pub fn format_control_result_as_links_notation(
393 action: ControlAction,
394 identifier: &str,
395 record: &ExecutionRecord,
396 method: &str,
397 process_ids: Option<&Value>,
398 message: &str,
399) -> String {
400 let backend = record
401 .options
402 .get("isolated")
403 .and_then(|value| value.as_str())
404 .unwrap_or("unknown");
405 let session_name = record
406 .options
407 .get("sessionName")
408 .and_then(|value| value.as_str())
409 .unwrap_or("");
410 let status = match action {
411 ControlAction::Stop => "signal-sent",
412 ControlAction::Terminate => "terminated",
413 };
414
415 let mut lines = vec![
416 "executionControl".to_string(),
417 format!(" action {}", escape_for_links_notation(action.as_str())),
418 format!(" identifier {}", escape_for_links_notation(identifier)),
419 format!(" uuid {}", escape_for_links_notation(&record.uuid)),
420 format!(" status {}", escape_for_links_notation(status)),
421 format!(" backend {}", escape_for_links_notation(backend)),
422 format!(" sessionName {}", escape_for_links_notation(session_name)),
423 format!(" method {}", escape_for_links_notation(method)),
424 ];
425
426 if let Some(process_ids) = process_ids {
427 append_links_value(&mut lines, "processIds", process_ids, 2);
428 }
429
430 lines.push(format!(" message {}", escape_for_links_notation(message)));
431 lines.join("\n")
432}
433
434pub fn control_execution(
435 store: Option<&ExecutionStore>,
436 identifier: &str,
437 action: ControlAction,
438) -> ExecutionControlResult {
439 control_execution_with_runner(store, identifier, action, &SystemCommandRunner)
440}
441
442pub fn control_execution_with_runner<R: CommandRunner>(
443 store: Option<&ExecutionStore>,
444 identifier: &str,
445 action: ControlAction,
446 runner: &R,
447) -> ExecutionControlResult {
448 let Some(store) = store else {
449 return ExecutionControlResult {
450 success: false,
451 output: None,
452 error: Some("Execution tracking is disabled.".to_string()),
453 };
454 };
455
456 let Some(record) = store.get(identifier) else {
457 return ExecutionControlResult {
458 success: false,
459 output: None,
460 error: Some(format!(
461 "No execution found with UUID or session name: {}",
462 identifier
463 )),
464 };
465 };
466
467 let control = match get_control_command(&record, action) {
468 Ok(command) => command,
469 Err(error) => {
470 return ExecutionControlResult {
471 success: false,
472 output: None,
473 error: Some(error),
474 }
475 }
476 };
477
478 let result = runner.run(&control.command, &control.args);
479 if !result.success {
480 let backend = record
481 .options
482 .get("isolated")
483 .and_then(|value| value.as_str())
484 .unwrap_or("unknown");
485 let session_name = record
486 .options
487 .get("sessionName")
488 .and_then(|value| value.as_str())
489 .unwrap_or("");
490 let detail = if !result.stderr.is_empty() {
491 result.stderr
492 } else if let Some(error) = result.error {
493 error
494 } else {
495 format!("exit code {}", result.status.unwrap_or(-1))
496 };
497
498 return ExecutionControlResult {
499 success: false,
500 output: None,
501 error: Some(format!(
502 "Failed to {} {} session \"{}\": {}",
503 action.as_str(),
504 backend,
505 session_name,
506 detail
507 )),
508 };
509 }
510
511 let process_ids = collect_process_ids_with_runner(&record, runner);
512 let output = format_control_result_as_links_notation(
513 action,
514 identifier,
515 &record,
516 &control.method,
517 process_ids.as_ref(),
518 &control.message,
519 );
520
521 ExecutionControlResult {
522 success: true,
523 output: Some(output),
524 error: None,
525 }
526}