Skip to main content

aft/commands/
bash_orchestrate.rs

1use std::path::Path;
2use std::time::{Duration, Instant};
3
4use serde::Deserialize;
5use serde_json::{json, Value};
6
7use crate::bash_background::registry::BgTaskSnapshot;
8use crate::bash_background::BgTaskStatus;
9use crate::context::AppContext;
10use crate::protocol::{RawRequest, Response};
11use crate::response_finalize::{DispatchOutcome, PendingResponse, PendingResponsePoll};
12
13const TEST_FOREGROUND_WAIT_ENV: &str = "AFT_TEST_FOREGROUND_WAIT_MS";
14const DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS: u64 = 30 * 60 * 1000;
15
16#[derive(Debug, Default, Deserialize)]
17#[serde(default)]
18struct BashOrchestrateParams {
19    foreground_orchestrate: bool,
20    block_to_completion: bool,
21    wait: bool,
22    background: bool,
23    pty: bool,
24    timeout: Option<u64>,
25}
26
27/// Port of `packages/aft-bridge/src/bash-format.ts` `formatForegroundResult` (lines 8-25).
28pub fn format_foreground_result(snapshot: &BgTaskSnapshot) -> String {
29    let mut rendered = snapshot.output_preview.clone();
30    if snapshot.output_truncated {
31        if let Some(output_path) = snapshot.output_path.as_deref() {
32            rendered.push_str(&format!(
33                "\n[output truncated; full output at {output_path}]"
34            ));
35        }
36    }
37    if snapshot.info.status == BgTaskStatus::TimedOut {
38        rendered.push_str("\n[command timed out]");
39    }
40    if let Some(exit) = snapshot.exit_code.filter(|exit| *exit != 0) {
41        rendered.push_str(&format!("\n[exit code: {exit}]"));
42    }
43    rendered
44}
45
46/// Port of `packages/aft-bridge/src/bash-format.ts` `formatSeconds` (lines 3-6).
47pub fn format_seconds(ms: u64) -> String {
48    let mut seconds = format!("{:.1}", ms as f64 / 1000.0);
49    if seconds.ends_with(".0") {
50        seconds.truncate(seconds.len() - 2);
51    }
52    format!("{seconds}s")
53}
54
55fn format_background_handoff_tail(task_id: &str) -> String {
56    format!(
57        "{task_id}. A completion reminder will be delivered automatically; use bash_status({{ taskId: \"{task_id}\" }}) to inspect output or bash_kill({{ taskId: \"{task_id}\" }}) to terminate."
58    )
59}
60
61/// Port of OpenCode `packages/opencode-plugin/src/tools/bash.ts` `formatPromotionMessage` (lines 603-614).
62pub fn format_promotion_message(
63    task_id: &str,
64    timeout: Option<u64>,
65    wait_window_ms: u64,
66) -> String {
67    let waited = timeout
68        .map(|timeout| timeout.min(wait_window_ms))
69        .unwrap_or(wait_window_ms);
70    format!(
71        "Foreground bash didn't finish within {} and was promoted to background: {}",
72        format_seconds(waited),
73        format_background_handoff_tail(task_id)
74    )
75}
76
77pub fn format_wait_detach_message(task_id: &str) -> String {
78    format!(
79        "Foreground bash is running in background as {}\nDetached because a user message arrived.",
80        format_background_handoff_tail(task_id)
81    )
82}
83
84/// Port of OpenCode `packages/opencode-plugin/src/tools/bash.ts` `formatBackgroundLaunch` (lines 593-601).
85pub fn format_background_launch(task_id: &str, pty: bool) -> String {
86    if pty {
87        return format!(
88            "PTY task started: {task_id}. Use bash_status({{ taskId: \"{task_id}\", outputMode: \"screen\" }}) to see the visible terminal, bash_write({{ taskId: \"{task_id}\", input: ... }}) to send keystrokes. A completion reminder fires automatically when the task exits."
89        );
90    }
91    format!(
92        "Background task started: {task_id}. A completion reminder will be delivered automatically; don't poll bash_status."
93    )
94}
95
96pub fn foreground_orchestrate_enabled(req: &RawRequest) -> bool {
97    parse_params(req)
98        .map(|params| params.foreground_orchestrate)
99        .unwrap_or(false)
100}
101
102pub fn build_bash_outcome(
103    req: &RawRequest,
104    ctx: &AppContext,
105    spawn_response: Response,
106) -> DispatchOutcome {
107    if !spawn_response.success {
108        return DispatchOutcome::Immediate(spawn_response);
109    }
110
111    let params = parse_params(req).unwrap_or_default();
112    let Some(task_id) = spawn_response
113        .data
114        .get("task_id")
115        .and_then(Value::as_str)
116        .map(str::to_owned)
117    else {
118        return DispatchOutcome::Immediate(spawn_response);
119    };
120    if spawn_response.data.get("status").and_then(Value::as_str) != Some("running") {
121        return DispatchOutcome::Immediate(spawn_response);
122    }
123
124    let mode = spawn_response
125        .data
126        .get("mode")
127        .and_then(Value::as_str)
128        .unwrap_or("pipes");
129    let is_pty = mode == "pty" || params.pty;
130    if is_pty || params.background {
131        return DispatchOutcome::Immediate(background_launch_response(&req.id, &task_id, is_pty));
132    }
133
134    let request_id = req.id.clone();
135    let session_id = req.session().to_string();
136    let attach_command = "bash".to_string();
137    let detach_on_user_message = params.wait;
138    ctx.bash_background()
139        .register_foreground_task(&session_id, &task_id);
140    if detach_on_user_message {
141        ctx.bash_background()
142            .begin_wait_mode_session(&session_id, &task_id);
143    }
144    let wait_window_ms = select_foreground_wait_window_ms(
145        ctx.config().foreground_wait_window_ms,
146        params.timeout,
147        params.wait,
148    );
149    let deadline = Instant::now() + Duration::from_millis(wait_window_ms);
150    let block_to_completion = params.block_to_completion || params.wait;
151    let timeout = params.timeout;
152    let storage_dir = crate::bash_background::storage_dir(ctx.config().storage_dir.as_deref());
153    let project_root = ctx.config().project_root.clone();
154    let task_id_for_poll = task_id.clone();
155    let request_id_for_poll = request_id.clone();
156    let session_id_for_poll = session_id.clone();
157    let session_id_for_cleanup = session_id.clone();
158
159    let mut poll: PendingResponsePoll = Box::new(move |ctx| {
160        // Foreground polls only need task state. Terminal snapshots still
161        // return cached output.
162        let response = if let Some(snapshot) = poll_bash_status(
163            ctx,
164            &task_id_for_poll,
165            &session_id_for_poll,
166            project_root.as_deref(),
167            &storage_dir,
168            0,
169        ) {
170            if snapshot.info.status.is_terminal() {
171                Some(foreground_result_response(&request_id_for_poll, snapshot))
172            } else if detach_on_user_message
173                && ctx
174                    .bash_background()
175                    .take_wait_mode_detach(&session_id_for_poll)
176            {
177                Some(detach_wait_mode_bash(
178                    ctx,
179                    &task_id_for_poll,
180                    &session_id_for_poll,
181                    &request_id_for_poll,
182                ))
183            } else {
184                match decide_bash_step(
185                    snapshot,
186                    deadline,
187                    block_to_completion,
188                    Instant::now(),
189                    &request_id_for_poll,
190                ) {
191                    BashStep::Done(response) => Some(response),
192                    BashStep::Promote => Some(promote_bash(
193                        ctx,
194                        &task_id_for_poll,
195                        &session_id_for_poll,
196                        timeout,
197                        wait_window_ms,
198                        &request_id_for_poll,
199                    )),
200                    BashStep::Wait => None,
201                }
202            }
203        } else {
204            Some(task_not_found_response(
205                &request_id_for_poll,
206                &task_id_for_poll,
207            ))
208        };
209
210        if response.is_some() {
211            if detach_on_user_message {
212                ctx.bash_background()
213                    .end_wait_mode_session(&session_id_for_cleanup, &task_id_for_poll);
214            } else {
215                ctx.bash_background()
216                    .unregister_foreground_task(&session_id_for_cleanup, &task_id_for_poll);
217            }
218        }
219        response
220    });
221
222    if let Some(response) = poll(ctx) {
223        return DispatchOutcome::Immediate(response);
224    }
225
226    DispatchOutcome::Deferred(PendingResponse {
227        request_id,
228        session_id,
229        attach_command,
230        poll,
231    })
232}
233
234pub(crate) fn poll_bash_status(
235    ctx: &AppContext,
236    task_id: &str,
237    session_id: &str,
238    project_root: Option<&Path>,
239    storage_dir: &Path,
240    preview_bytes: usize,
241) -> Option<BgTaskSnapshot> {
242    ctx.bash_background().status(
243        task_id,
244        session_id,
245        project_root,
246        Some(storage_dir),
247        preview_bytes,
248    )
249}
250
251pub(crate) enum BashStep {
252    Done(Response),
253    Promote,
254    Wait,
255}
256
257pub(crate) fn task_not_found_response(request_id: &str, task_id: &str) -> Response {
258    Response::error(
259        request_id,
260        "task_not_found",
261        format!("background task not found: {task_id}"),
262    )
263}
264
265pub(crate) fn decide_bash_step(
266    snapshot: BgTaskSnapshot,
267    deadline: Instant,
268    block_to_completion: bool,
269    now: Instant,
270    request_id: &str,
271) -> BashStep {
272    if snapshot.info.status.is_terminal() {
273        BashStep::Done(foreground_result_response(request_id, snapshot))
274    } else if !block_to_completion && now >= deadline {
275        BashStep::Promote
276    } else {
277        BashStep::Wait
278    }
279}
280
281pub(crate) fn promote_bash(
282    ctx: &AppContext,
283    task_id: &str,
284    session_id: &str,
285    timeout: Option<u64>,
286    wait_window_ms: u64,
287    request_id: &str,
288) -> Response {
289    match ctx.bash_background().promote(task_id, session_id) {
290        Ok(_) => promotion_response(request_id, task_id, timeout, wait_window_ms),
291        Err(message) if message.contains("not found") => {
292            Response::error(request_id, "task_not_found", message)
293        }
294        Err(message) => Response::error(request_id, "execution_failed", message),
295    }
296}
297
298pub(crate) fn detach_wait_mode_bash(
299    ctx: &AppContext,
300    task_id: &str,
301    session_id: &str,
302    request_id: &str,
303) -> Response {
304    match ctx.bash_background().promote(task_id, session_id) {
305        Ok(_) => wait_detach_response(request_id, task_id),
306        Err(message) if message.contains("not found") => {
307            Response::error(request_id, "task_not_found", message)
308        }
309        Err(message) => Response::error(request_id, "execution_failed", message),
310    }
311}
312
313fn foreground_result_response(request_id: &str, snapshot: BgTaskSnapshot) -> Response {
314    let output = format_foreground_result(&snapshot);
315    if snapshot.sandbox_native
316        && snapshot.sandbox_unavailable
317        && snapshot.exit_code == Some(crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE)
318    {
319        return Response::error_with_data(
320            request_id,
321            "sandbox_unavailable",
322            "native sandbox failed before the command could run; set sandbox.enabled=false to disable native sandboxing",
323            json!({
324                "task_id": snapshot.info.task_id,
325                "exit_code": snapshot.exit_code,
326                "output": output,
327            }),
328        );
329    }
330    let timed_out = snapshot.info.status == BgTaskStatus::TimedOut;
331    Response::success(
332        request_id,
333        json!({
334            "output": output,
335            "task_id": snapshot.info.task_id,
336            "status": snapshot.info.status,
337            "mode": snapshot.info.mode,
338            "exit_code": snapshot.exit_code,
339            "output_preview": snapshot.output_preview,
340            "output_truncated": snapshot.output_truncated,
341            "truncated": snapshot.output_truncated,
342            "output_path": snapshot.output_path,
343            "timed_out": timed_out,
344            "duration_ms": snapshot.info.duration_ms,
345        }),
346    )
347}
348
349fn background_launch_response(request_id: &str, task_id: &str, is_pty: bool) -> Response {
350    Response::success(
351        request_id,
352        json!({
353            "output": format_background_launch(task_id, is_pty),
354            "task_id": task_id,
355            "status": "running",
356            "mode": if is_pty { "pty" } else { "pipes" },
357        }),
358    )
359}
360
361fn promotion_response(
362    request_id: &str,
363    task_id: &str,
364    timeout: Option<u64>,
365    wait_window_ms: u64,
366) -> Response {
367    Response::success(
368        request_id,
369        json!({
370            "output": format_promotion_message(task_id, timeout, wait_window_ms),
371            "task_id": task_id,
372            "status": "running",
373        }),
374    )
375}
376
377fn wait_detach_response(request_id: &str, task_id: &str) -> Response {
378    Response::success(
379        request_id,
380        json!({
381            "output": format_wait_detach_message(task_id),
382            "task_id": task_id,
383            "status": "running",
384        }),
385    )
386}
387
388fn parse_params(req: &RawRequest) -> Option<BashOrchestrateParams> {
389    let raw_params = req
390        .params
391        .get("params")
392        .cloned()
393        .unwrap_or_else(|| req.params.clone());
394    serde_json::from_value::<BashOrchestrateParams>(raw_params).ok()
395}
396
397pub(crate) fn resolve_foreground_wait_window_ms(configured: u64) -> u64 {
398    std::env::var(TEST_FOREGROUND_WAIT_ENV)
399        .ok()
400        .and_then(|raw| raw.parse::<u64>().ok())
401        .unwrap_or(configured)
402}
403
404pub(crate) fn select_foreground_wait_window_ms(
405    configured: u64,
406    timeout: Option<u64>,
407    wait: bool,
408) -> u64 {
409    if wait {
410        timeout.unwrap_or(DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS)
411    } else {
412        resolve_foreground_wait_window_ms(configured)
413    }
414}
415
416#[cfg(test)]
417mod tests {
418    use super::*;
419    use crate::bash_background::persistence::BgMode;
420    use crate::bash_background::registry::BgTaskSnapshot;
421    use crate::bash_background::BgTaskInfo;
422
423    fn snapshot(
424        output_preview: &str,
425        output_truncated: bool,
426        output_path: Option<&str>,
427        status: BgTaskStatus,
428        exit_code: Option<i32>,
429    ) -> BgTaskSnapshot {
430        BgTaskSnapshot {
431            info: BgTaskInfo {
432                task_id: "bash-test".to_string(),
433                status,
434                command: "echo test".to_string(),
435                mode: BgMode::Pipes,
436                started_at: 0,
437                duration_ms: Some(1),
438                status_reason: None,
439            },
440            exit_code,
441            child_pid: None,
442            workdir: "/tmp".to_string(),
443            output_preview: output_preview.to_string(),
444            output_truncated,
445            output_path: output_path.map(str::to_string),
446            stderr_path: None,
447            pty_rows: None,
448            pty_cols: None,
449            pty_screen: None,
450            scanner_report: Vec::new(),
451            sandbox_native: false,
452            sandbox_unavailable: false,
453        }
454    }
455
456    #[test]
457    fn native_launcher_exit_78_is_a_structured_sandbox_error() {
458        let mut snapshot = snapshot(
459            "sandbox_unavailable: backend failed",
460            false,
461            None,
462            BgTaskStatus::Failed,
463            Some(crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE),
464        );
465        snapshot.sandbox_native = true;
466        snapshot.sandbox_unavailable = true;
467
468        let response = foreground_result_response("sandbox-failed", snapshot);
469        assert!(!response.success);
470        assert_eq!(
471            response
472                .data
473                .get("code")
474                .and_then(serde_json::Value::as_str),
475            Some("sandbox_unavailable")
476        );
477        assert!(response
478            .data
479            .get("message")
480            .and_then(serde_json::Value::as_str)
481            .is_some_and(|message| message.contains("sandbox.enabled=false")));
482    }
483
484    #[test]
485    fn native_command_exit_78_is_not_misreported_as_launcher_failure() {
486        let mut snapshot = snapshot(
487            "command selected exit 78",
488            false,
489            None,
490            BgTaskStatus::Failed,
491            Some(crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE),
492        );
493        snapshot.sandbox_native = true;
494
495        let response = foreground_result_response("command-exit-78", snapshot);
496        assert!(response.success);
497        assert_eq!(
498            response
499                .data
500                .get("exit_code")
501                .and_then(serde_json::Value::as_i64),
502            Some(i64::from(
503                crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE
504            ))
505        );
506    }
507
508    #[test]
509    fn decide_bash_step_returns_done_for_terminal_snapshot_even_at_deadline() {
510        let snapshot = snapshot("done", false, None, BgTaskStatus::Completed, Some(0));
511        let now = Instant::now();
512
513        match decide_bash_step(snapshot, now, false, now, "req-terminal") {
514            BashStep::Done(response) => {
515                assert_eq!(response.id, "req-terminal");
516                assert!(response.success);
517                assert_eq!(response.data["status"], json!("completed"));
518                assert_eq!(response.data["output"], json!("done"));
519            }
520            BashStep::Promote => panic!("terminal snapshot should not promote"),
521            BashStep::Wait => panic!("terminal snapshot should not wait"),
522        }
523    }
524
525    #[test]
526    fn decide_bash_step_promotes_at_deadline_when_not_blocking() {
527        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
528        let now = Instant::now();
529
530        match decide_bash_step(snapshot, now, false, now, "req-promote") {
531            BashStep::Promote => {}
532            BashStep::Done(_) => panic!("running snapshot should not finish"),
533            BashStep::Wait => panic!("deadline should promote when not blocking"),
534        }
535    }
536
537    #[test]
538    fn decide_bash_step_waits_before_deadline() {
539        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
540        let now = Instant::now();
541
542        match decide_bash_step(
543            snapshot,
544            now + Duration::from_millis(1),
545            false,
546            now,
547            "req-wait",
548        ) {
549            BashStep::Wait => {}
550            BashStep::Done(_) => panic!("running snapshot should not finish"),
551            BashStep::Promote => panic!("snapshot should wait before the deadline"),
552        }
553    }
554
555    #[test]
556    fn decide_bash_step_never_promotes_when_blocking_to_completion() {
557        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
558        let now = Instant::now();
559
560        match decide_bash_step(snapshot, now, true, now, "req-block") {
561            BashStep::Wait => {}
562            BashStep::Done(_) => panic!("running snapshot should not finish"),
563            BashStep::Promote => panic!("block_to_completion should suppress promotion"),
564        }
565    }
566
567    #[test]
568    fn select_foreground_wait_window_uses_timeout_budget_for_wait_true() {
569        assert_eq!(
570            select_foreground_wait_window_ms(8_000, Some(250), true),
571            250
572        );
573        assert_eq!(
574            select_foreground_wait_window_ms(8_000, None, true),
575            DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS
576        );
577    }
578
579    #[test]
580    fn foreground_result_format_matches_typescript_order() {
581        let snapshot = snapshot(
582            "hello",
583            true,
584            Some("/tmp/aft-output.txt"),
585            BgTaskStatus::TimedOut,
586            Some(124),
587        );
588
589        assert_eq!(
590            format_foreground_result(&snapshot),
591            "hello\n[output truncated; full output at /tmp/aft-output.txt]\n[command timed out]\n[exit code: 124]"
592        );
593    }
594
595    #[test]
596    fn format_seconds_strips_integer_decimal_and_keeps_tenths() {
597        assert_eq!(format_seconds(8_000), "8s");
598        assert_eq!(format_seconds(5_500), "5.5s");
599        assert_eq!(format_seconds(14_999), "15s");
600    }
601
602    #[test]
603    fn promotion_message_matches_opencode_copy() {
604        assert_eq!(
605            format_promotion_message("bash-123", Some(5_500), 8_000),
606            "Foreground bash didn't finish within 5.5s and was promoted to background: bash-123. A completion reminder will be delivered automatically; use bash_status({ taskId: \"bash-123\" }) to inspect output or bash_kill({ taskId: \"bash-123\" }) to terminate."
607        );
608    }
609
610    #[test]
611    fn wait_detach_message_mentions_user_message() {
612        assert_eq!(
613            format_wait_detach_message("bash-123"),
614            "Foreground bash is running in background as bash-123. A completion reminder will be delivered automatically; use bash_status({ taskId: \"bash-123\" }) to inspect output or bash_kill({ taskId: \"bash-123\" }) to terminate.\nDetached because a user message arrived."
615        );
616    }
617
618    #[test]
619    fn background_launch_messages_match_opencode_copy() {
620        assert_eq!(
621            format_background_launch("bash-bg", false),
622            "Background task started: bash-bg. A completion reminder will be delivered automatically; don't poll bash_status."
623        );
624        assert_eq!(
625            format_background_launch("bash-pty", true),
626            "PTY task started: bash-pty. Use bash_status({ taskId: \"bash-pty\", outputMode: \"screen\" }) to see the visible terminal, bash_write({ taskId: \"bash-pty\", input: ... }) to send keystrokes. A completion reminder fires automatically when the task exits."
627        );
628    }
629}