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    if detach_on_user_message {
139        ctx.bash_background().begin_wait_mode_session(&session_id);
140    }
141    let wait_window_ms = select_foreground_wait_window_ms(
142        ctx.config().foreground_wait_window_ms,
143        params.timeout,
144        params.wait,
145    );
146    let deadline = Instant::now() + Duration::from_millis(wait_window_ms);
147    let block_to_completion = params.block_to_completion || params.wait;
148    let timeout = params.timeout;
149    let storage_dir = crate::bash_background::storage_dir(ctx.config().storage_dir.as_deref());
150    let project_root = ctx.config().project_root.clone();
151    let task_id_for_poll = task_id.clone();
152    let request_id_for_poll = request_id.clone();
153    let session_id_for_poll = session_id.clone();
154    let session_id_for_cleanup = session_id.clone();
155
156    let mut poll: PendingResponsePoll = Box::new(move |ctx| {
157        // Foreground polls only need task state. Terminal snapshots still
158        // return cached output.
159        let response = if let Some(snapshot) = poll_bash_status(
160            ctx,
161            &task_id_for_poll,
162            &session_id_for_poll,
163            project_root.as_deref(),
164            &storage_dir,
165            0,
166        ) {
167            if snapshot.info.status.is_terminal() {
168                Some(foreground_result_response(&request_id_for_poll, snapshot))
169            } else if detach_on_user_message
170                && ctx
171                    .bash_background()
172                    .take_wait_mode_detach(&session_id_for_poll)
173            {
174                Some(detach_wait_mode_bash(
175                    ctx,
176                    &task_id_for_poll,
177                    &session_id_for_poll,
178                    &request_id_for_poll,
179                ))
180            } else {
181                match decide_bash_step(
182                    snapshot,
183                    deadline,
184                    block_to_completion,
185                    Instant::now(),
186                    &request_id_for_poll,
187                ) {
188                    BashStep::Done(response) => Some(response),
189                    BashStep::Promote => Some(promote_bash(
190                        ctx,
191                        &task_id_for_poll,
192                        &session_id_for_poll,
193                        timeout,
194                        wait_window_ms,
195                        &request_id_for_poll,
196                    )),
197                    BashStep::Wait => None,
198                }
199            }
200        } else {
201            Some(task_not_found_response(
202                &request_id_for_poll,
203                &task_id_for_poll,
204            ))
205        };
206
207        if response.is_some() && detach_on_user_message {
208            ctx.bash_background()
209                .end_wait_mode_session(&session_id_for_cleanup);
210        }
211        response
212    });
213
214    if let Some(response) = poll(ctx) {
215        return DispatchOutcome::Immediate(response);
216    }
217
218    DispatchOutcome::Deferred(PendingResponse {
219        request_id,
220        session_id,
221        attach_command,
222        poll,
223    })
224}
225
226pub(crate) fn poll_bash_status(
227    ctx: &AppContext,
228    task_id: &str,
229    session_id: &str,
230    project_root: Option<&Path>,
231    storage_dir: &Path,
232    preview_bytes: usize,
233) -> Option<BgTaskSnapshot> {
234    ctx.bash_background().status(
235        task_id,
236        session_id,
237        project_root,
238        Some(storage_dir),
239        preview_bytes,
240    )
241}
242
243pub(crate) enum BashStep {
244    Done(Response),
245    Promote,
246    Wait,
247}
248
249pub(crate) fn task_not_found_response(request_id: &str, task_id: &str) -> Response {
250    Response::error(
251        request_id,
252        "task_not_found",
253        format!("background task not found: {task_id}"),
254    )
255}
256
257pub(crate) fn decide_bash_step(
258    snapshot: BgTaskSnapshot,
259    deadline: Instant,
260    block_to_completion: bool,
261    now: Instant,
262    request_id: &str,
263) -> BashStep {
264    if snapshot.info.status.is_terminal() {
265        BashStep::Done(foreground_result_response(request_id, snapshot))
266    } else if !block_to_completion && now >= deadline {
267        BashStep::Promote
268    } else {
269        BashStep::Wait
270    }
271}
272
273pub(crate) fn promote_bash(
274    ctx: &AppContext,
275    task_id: &str,
276    session_id: &str,
277    timeout: Option<u64>,
278    wait_window_ms: u64,
279    request_id: &str,
280) -> Response {
281    match ctx.bash_background().promote(task_id, session_id) {
282        Ok(_) => promotion_response(request_id, task_id, timeout, wait_window_ms),
283        Err(message) if message.contains("not found") => {
284            Response::error(request_id, "task_not_found", message)
285        }
286        Err(message) => Response::error(request_id, "execution_failed", message),
287    }
288}
289
290pub(crate) fn detach_wait_mode_bash(
291    ctx: &AppContext,
292    task_id: &str,
293    session_id: &str,
294    request_id: &str,
295) -> Response {
296    match ctx.bash_background().promote(task_id, session_id) {
297        Ok(_) => wait_detach_response(request_id, task_id),
298        Err(message) if message.contains("not found") => {
299            Response::error(request_id, "task_not_found", message)
300        }
301        Err(message) => Response::error(request_id, "execution_failed", message),
302    }
303}
304
305fn foreground_result_response(request_id: &str, snapshot: BgTaskSnapshot) -> Response {
306    let output = format_foreground_result(&snapshot);
307    if snapshot.sandbox_native
308        && snapshot.sandbox_unavailable
309        && snapshot.exit_code == Some(crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE)
310    {
311        return Response::error_with_data(
312            request_id,
313            "sandbox_unavailable",
314            "native sandbox failed before the command could run; set sandbox.enabled=false to disable native sandboxing",
315            json!({
316                "task_id": snapshot.info.task_id,
317                "exit_code": snapshot.exit_code,
318                "output": output,
319            }),
320        );
321    }
322    let timed_out = snapshot.info.status == BgTaskStatus::TimedOut;
323    Response::success(
324        request_id,
325        json!({
326            "output": output,
327            "task_id": snapshot.info.task_id,
328            "status": snapshot.info.status,
329            "mode": snapshot.info.mode,
330            "exit_code": snapshot.exit_code,
331            "output_preview": snapshot.output_preview,
332            "output_truncated": snapshot.output_truncated,
333            "truncated": snapshot.output_truncated,
334            "output_path": snapshot.output_path,
335            "timed_out": timed_out,
336            "duration_ms": snapshot.info.duration_ms,
337        }),
338    )
339}
340
341fn background_launch_response(request_id: &str, task_id: &str, is_pty: bool) -> Response {
342    Response::success(
343        request_id,
344        json!({
345            "output": format_background_launch(task_id, is_pty),
346            "task_id": task_id,
347            "status": "running",
348            "mode": if is_pty { "pty" } else { "pipes" },
349        }),
350    )
351}
352
353fn promotion_response(
354    request_id: &str,
355    task_id: &str,
356    timeout: Option<u64>,
357    wait_window_ms: u64,
358) -> Response {
359    Response::success(
360        request_id,
361        json!({
362            "output": format_promotion_message(task_id, timeout, wait_window_ms),
363            "task_id": task_id,
364            "status": "running",
365        }),
366    )
367}
368
369fn wait_detach_response(request_id: &str, task_id: &str) -> Response {
370    Response::success(
371        request_id,
372        json!({
373            "output": format_wait_detach_message(task_id),
374            "task_id": task_id,
375            "status": "running",
376        }),
377    )
378}
379
380fn parse_params(req: &RawRequest) -> Option<BashOrchestrateParams> {
381    let raw_params = req
382        .params
383        .get("params")
384        .cloned()
385        .unwrap_or_else(|| req.params.clone());
386    serde_json::from_value::<BashOrchestrateParams>(raw_params).ok()
387}
388
389pub(crate) fn resolve_foreground_wait_window_ms(configured: u64) -> u64 {
390    std::env::var(TEST_FOREGROUND_WAIT_ENV)
391        .ok()
392        .and_then(|raw| raw.parse::<u64>().ok())
393        .unwrap_or(configured)
394}
395
396pub(crate) fn select_foreground_wait_window_ms(
397    configured: u64,
398    timeout: Option<u64>,
399    wait: bool,
400) -> u64 {
401    if wait {
402        timeout.unwrap_or(DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS)
403    } else {
404        resolve_foreground_wait_window_ms(configured)
405    }
406}
407
408#[cfg(test)]
409mod tests {
410    use super::*;
411    use crate::bash_background::persistence::BgMode;
412    use crate::bash_background::registry::BgTaskSnapshot;
413    use crate::bash_background::BgTaskInfo;
414
415    fn snapshot(
416        output_preview: &str,
417        output_truncated: bool,
418        output_path: Option<&str>,
419        status: BgTaskStatus,
420        exit_code: Option<i32>,
421    ) -> BgTaskSnapshot {
422        BgTaskSnapshot {
423            info: BgTaskInfo {
424                task_id: "bash-test".to_string(),
425                status,
426                command: "echo test".to_string(),
427                mode: BgMode::Pipes,
428                started_at: 0,
429                duration_ms: Some(1),
430            },
431            exit_code,
432            child_pid: None,
433            workdir: "/tmp".to_string(),
434            output_preview: output_preview.to_string(),
435            output_truncated,
436            output_path: output_path.map(str::to_string),
437            stderr_path: None,
438            pty_rows: None,
439            pty_cols: None,
440            pty_screen: None,
441            scanner_report: Vec::new(),
442            sandbox_native: false,
443            sandbox_unavailable: false,
444        }
445    }
446
447    #[test]
448    fn native_launcher_exit_78_is_a_structured_sandbox_error() {
449        let mut snapshot = snapshot(
450            "sandbox_unavailable: backend failed",
451            false,
452            None,
453            BgTaskStatus::Failed,
454            Some(crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE),
455        );
456        snapshot.sandbox_native = true;
457        snapshot.sandbox_unavailable = true;
458
459        let response = foreground_result_response("sandbox-failed", snapshot);
460        assert!(!response.success);
461        assert_eq!(
462            response
463                .data
464                .get("code")
465                .and_then(serde_json::Value::as_str),
466            Some("sandbox_unavailable")
467        );
468        assert!(response
469            .data
470            .get("message")
471            .and_then(serde_json::Value::as_str)
472            .is_some_and(|message| message.contains("sandbox.enabled=false")));
473    }
474
475    #[test]
476    fn native_command_exit_78_is_not_misreported_as_launcher_failure() {
477        let mut snapshot = snapshot(
478            "command selected exit 78",
479            false,
480            None,
481            BgTaskStatus::Failed,
482            Some(crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE),
483        );
484        snapshot.sandbox_native = true;
485
486        let response = foreground_result_response("command-exit-78", snapshot);
487        assert!(response.success);
488        assert_eq!(
489            response
490                .data
491                .get("exit_code")
492                .and_then(serde_json::Value::as_i64),
493            Some(i64::from(
494                crate::sandbox_spawn::SANDBOX_UNAVAILABLE_EXIT_CODE
495            ))
496        );
497    }
498
499    #[test]
500    fn decide_bash_step_returns_done_for_terminal_snapshot_even_at_deadline() {
501        let snapshot = snapshot("done", false, None, BgTaskStatus::Completed, Some(0));
502        let now = Instant::now();
503
504        match decide_bash_step(snapshot, now, false, now, "req-terminal") {
505            BashStep::Done(response) => {
506                assert_eq!(response.id, "req-terminal");
507                assert!(response.success);
508                assert_eq!(response.data["status"], json!("completed"));
509                assert_eq!(response.data["output"], json!("done"));
510            }
511            BashStep::Promote => panic!("terminal snapshot should not promote"),
512            BashStep::Wait => panic!("terminal snapshot should not wait"),
513        }
514    }
515
516    #[test]
517    fn decide_bash_step_promotes_at_deadline_when_not_blocking() {
518        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
519        let now = Instant::now();
520
521        match decide_bash_step(snapshot, now, false, now, "req-promote") {
522            BashStep::Promote => {}
523            BashStep::Done(_) => panic!("running snapshot should not finish"),
524            BashStep::Wait => panic!("deadline should promote when not blocking"),
525        }
526    }
527
528    #[test]
529    fn decide_bash_step_waits_before_deadline() {
530        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
531        let now = Instant::now();
532
533        match decide_bash_step(
534            snapshot,
535            now + Duration::from_millis(1),
536            false,
537            now,
538            "req-wait",
539        ) {
540            BashStep::Wait => {}
541            BashStep::Done(_) => panic!("running snapshot should not finish"),
542            BashStep::Promote => panic!("snapshot should wait before the deadline"),
543        }
544    }
545
546    #[test]
547    fn decide_bash_step_never_promotes_when_blocking_to_completion() {
548        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
549        let now = Instant::now();
550
551        match decide_bash_step(snapshot, now, true, now, "req-block") {
552            BashStep::Wait => {}
553            BashStep::Done(_) => panic!("running snapshot should not finish"),
554            BashStep::Promote => panic!("block_to_completion should suppress promotion"),
555        }
556    }
557
558    #[test]
559    fn select_foreground_wait_window_uses_timeout_budget_for_wait_true() {
560        assert_eq!(
561            select_foreground_wait_window_ms(8_000, Some(250), true),
562            250
563        );
564        assert_eq!(
565            select_foreground_wait_window_ms(8_000, None, true),
566            DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS
567        );
568    }
569
570    #[test]
571    fn foreground_result_format_matches_typescript_order() {
572        let snapshot = snapshot(
573            "hello",
574            true,
575            Some("/tmp/aft-output.txt"),
576            BgTaskStatus::TimedOut,
577            Some(124),
578        );
579
580        assert_eq!(
581            format_foreground_result(&snapshot),
582            "hello\n[output truncated; full output at /tmp/aft-output.txt]\n[command timed out]\n[exit code: 124]"
583        );
584    }
585
586    #[test]
587    fn format_seconds_strips_integer_decimal_and_keeps_tenths() {
588        assert_eq!(format_seconds(8_000), "8s");
589        assert_eq!(format_seconds(5_500), "5.5s");
590        assert_eq!(format_seconds(14_999), "15s");
591    }
592
593    #[test]
594    fn promotion_message_matches_opencode_copy() {
595        assert_eq!(
596            format_promotion_message("bash-123", Some(5_500), 8_000),
597            "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."
598        );
599    }
600
601    #[test]
602    fn wait_detach_message_mentions_user_message() {
603        assert_eq!(
604            format_wait_detach_message("bash-123"),
605            "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."
606        );
607    }
608
609    #[test]
610    fn background_launch_messages_match_opencode_copy() {
611        assert_eq!(
612            format_background_launch("bash-bg", false),
613            "Background task started: bash-bg. A completion reminder will be delivered automatically; don't poll bash_status."
614        );
615        assert_eq!(
616            format_background_launch("bash-pty", true),
617            "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."
618        );
619    }
620}