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_array(lines: &mut Vec<String>, values: &[Value], indent: usize) {
371 let prefix = " ".repeat(indent);
372 if values.is_empty() {
373 lines.push(format!("{}()", prefix));
374 return;
375 }
376
377 lines.push(format!("{}(", prefix));
378 for value in values {
379 match value {
380 Value::Array(nested) => append_links_array(lines, nested, indent + 2),
381 Value::Object(map) => {
382 for (child_key, child_value) in map {
383 if !child_value.is_null() {
384 append_links_value(lines, child_key, child_value, indent + 2);
385 }
386 }
387 }
388 _ => lines.push(format!(
389 "{}{}",
390 " ".repeat(indent + 2),
391 format_value_for_links_notation(value)
392 )),
393 }
394 }
395 lines.push(format!("{})", prefix));
396}
397
398fn append_links_value(lines: &mut Vec<String>, key: &str, value: &Value, indent: usize) {
399 let prefix = " ".repeat(indent);
400 match value {
401 Value::Object(map) => {
402 lines.push(format!("{}{}", prefix, key));
403 for (child_key, child_value) in map {
404 if !child_value.is_null() {
405 append_links_value(lines, child_key, child_value, indent + 4);
406 }
407 }
408 }
409 Value::Array(values) => {
410 lines.push(format!("{}{}", prefix, key));
411 append_links_array(lines, values, indent + 2);
412 }
413 _ => {
414 lines.push(format!(
415 "{}{} {}",
416 prefix,
417 key,
418 format_value_for_links_notation(value)
419 ));
420 }
421 }
422}
423
424pub fn format_control_result_as_links_notation(
425 action: ControlAction,
426 identifier: &str,
427 record: &ExecutionRecord,
428 method: &str,
429 process_ids: Option<&Value>,
430 message: &str,
431) -> String {
432 let backend = record
433 .options
434 .get("isolated")
435 .and_then(|value| value.as_str())
436 .unwrap_or("unknown");
437 let session_name = record
438 .options
439 .get("sessionName")
440 .and_then(|value| value.as_str())
441 .unwrap_or("");
442 let status = match action {
443 ControlAction::Stop => "signal-sent",
444 ControlAction::Terminate => "terminated",
445 };
446
447 let mut lines = vec![
448 "executionControl".to_string(),
449 format!(" action {}", escape_for_links_notation(action.as_str())),
450 format!(" identifier {}", escape_for_links_notation(identifier)),
451 format!(" uuid {}", escape_for_links_notation(&record.uuid)),
452 format!(" status {}", escape_for_links_notation(status)),
453 format!(" backend {}", escape_for_links_notation(backend)),
454 format!(" sessionName {}", escape_for_links_notation(session_name)),
455 format!(" method {}", escape_for_links_notation(method)),
456 ];
457
458 if let Some(process_ids) = process_ids {
459 append_links_value(&mut lines, "processIds", process_ids, 2);
460 }
461
462 lines.push(format!(" message {}", escape_for_links_notation(message)));
463 lines.join("\n")
464}
465
466pub fn control_execution(
467 store: Option<&ExecutionStore>,
468 identifier: &str,
469 action: ControlAction,
470) -> ExecutionControlResult {
471 control_execution_with_runner(store, identifier, action, &SystemCommandRunner)
472}
473
474pub fn control_execution_with_runner<R: CommandRunner>(
475 store: Option<&ExecutionStore>,
476 identifier: &str,
477 action: ControlAction,
478 runner: &R,
479) -> ExecutionControlResult {
480 let Some(store) = store else {
481 return ExecutionControlResult {
482 success: false,
483 output: None,
484 error: Some("Execution tracking is disabled.".to_string()),
485 };
486 };
487
488 let Some(record) = store.get(identifier) else {
489 return ExecutionControlResult {
490 success: false,
491 output: None,
492 error: Some(format!(
493 "No execution found with UUID or session name: {}",
494 identifier
495 )),
496 };
497 };
498
499 let control = match get_control_command(&record, action) {
500 Ok(command) => command,
501 Err(error) => {
502 return ExecutionControlResult {
503 success: false,
504 output: None,
505 error: Some(error),
506 }
507 }
508 };
509
510 let result = runner.run(&control.command, &control.args);
511 if !result.success {
512 let backend = record
513 .options
514 .get("isolated")
515 .and_then(|value| value.as_str())
516 .unwrap_or("unknown");
517 let session_name = record
518 .options
519 .get("sessionName")
520 .and_then(|value| value.as_str())
521 .unwrap_or("");
522 let detail = if !result.stderr.is_empty() {
523 result.stderr
524 } else if let Some(error) = result.error {
525 error
526 } else {
527 format!("exit code {}", result.status.unwrap_or(-1))
528 };
529
530 return ExecutionControlResult {
531 success: false,
532 output: None,
533 error: Some(format!(
534 "Failed to {} {} session \"{}\": {}",
535 action.as_str(),
536 backend,
537 session_name,
538 detail
539 )),
540 };
541 }
542
543 let process_ids = collect_process_ids_with_runner(&record, runner);
544 let output = format_control_result_as_links_notation(
545 action,
546 identifier,
547 &record,
548 &control.method,
549 process_ids.as_ref(),
550 &control.message,
551 );
552
553 ExecutionControlResult {
554 success: true,
555 output: Some(output),
556 error: None,
557 }
558}