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::output::RUNNING_OUTPUT_PREVIEW_BYTES;
8use crate::bash_background::registry::BgTaskSnapshot;
9use crate::bash_background::BgTaskStatus;
10use crate::context::AppContext;
11use crate::protocol::{RawRequest, Response};
12use crate::response_finalize::{DispatchOutcome, PendingResponse, PendingResponsePoll};
13
14const TEST_FOREGROUND_WAIT_ENV: &str = "AFT_TEST_FOREGROUND_WAIT_MS";
15const DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS: u64 = 30 * 60 * 1000;
16
17#[derive(Debug, Default, Deserialize)]
18#[serde(default)]
19struct BashOrchestrateParams {
20    foreground_orchestrate: bool,
21    block_to_completion: bool,
22    wait: bool,
23    background: bool,
24    pty: bool,
25    timeout: Option<u64>,
26}
27
28/// Port of `packages/aft-bridge/src/bash-format.ts` `formatForegroundResult` (lines 8-25).
29pub fn format_foreground_result(snapshot: &BgTaskSnapshot) -> String {
30    let mut rendered = snapshot.output_preview.clone();
31    if snapshot.output_truncated {
32        if let Some(output_path) = snapshot.output_path.as_deref() {
33            rendered.push_str(&format!(
34                "\n[output truncated; full output at {output_path}]"
35            ));
36        }
37    }
38    if snapshot.info.status == BgTaskStatus::TimedOut {
39        rendered.push_str("\n[command timed out]");
40    }
41    if let Some(exit) = snapshot.exit_code.filter(|exit| *exit != 0) {
42        rendered.push_str(&format!("\n[exit code: {exit}]"));
43    }
44    rendered
45}
46
47/// Port of `packages/aft-bridge/src/bash-format.ts` `formatSeconds` (lines 3-6).
48pub fn format_seconds(ms: u64) -> String {
49    let mut seconds = format!("{:.1}", ms as f64 / 1000.0);
50    if seconds.ends_with(".0") {
51        seconds.truncate(seconds.len() - 2);
52    }
53    format!("{seconds}s")
54}
55
56/// Port of OpenCode `packages/opencode-plugin/src/tools/bash.ts` `formatPromotionMessage` (lines 603-614).
57pub fn format_promotion_message(
58    task_id: &str,
59    timeout: Option<u64>,
60    wait_window_ms: u64,
61) -> String {
62    let waited = timeout
63        .map(|timeout| timeout.min(wait_window_ms))
64        .unwrap_or(wait_window_ms);
65    format!(
66        "Foreground bash didn't finish within {} and was promoted to background: {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.",
67        format_seconds(waited)
68    )
69}
70
71/// Port of OpenCode `packages/opencode-plugin/src/tools/bash.ts` `formatBackgroundLaunch` (lines 593-601).
72pub fn format_background_launch(task_id: &str, pty: bool) -> String {
73    if pty {
74        return format!(
75            "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."
76        );
77    }
78    format!(
79        "Background task started: {task_id}. A completion reminder will be delivered automatically; don't poll bash_status."
80    )
81}
82
83pub fn foreground_orchestrate_enabled(req: &RawRequest) -> bool {
84    parse_params(req)
85        .map(|params| params.foreground_orchestrate)
86        .unwrap_or(false)
87}
88
89pub fn build_bash_outcome(
90    req: &RawRequest,
91    ctx: &AppContext,
92    spawn_response: Response,
93) -> DispatchOutcome {
94    if !spawn_response.success {
95        return DispatchOutcome::Immediate(spawn_response);
96    }
97
98    let params = parse_params(req).unwrap_or_default();
99    let Some(task_id) = spawn_response
100        .data
101        .get("task_id")
102        .and_then(Value::as_str)
103        .map(str::to_owned)
104    else {
105        return DispatchOutcome::Immediate(spawn_response);
106    };
107    if spawn_response.data.get("status").and_then(Value::as_str) != Some("running") {
108        return DispatchOutcome::Immediate(spawn_response);
109    }
110
111    let mode = spawn_response
112        .data
113        .get("mode")
114        .and_then(Value::as_str)
115        .unwrap_or("pipes");
116    let is_pty = mode == "pty" || params.pty;
117    if is_pty || params.background {
118        return DispatchOutcome::Immediate(background_launch_response(&req.id, &task_id, is_pty));
119    }
120
121    let request_id = req.id.clone();
122    let session_id = req.session().to_string();
123    let attach_command = "bash".to_string();
124    let wait_window_ms = select_foreground_wait_window_ms(
125        ctx.config().foreground_wait_window_ms,
126        params.timeout,
127        params.wait,
128    );
129    let deadline = Instant::now() + Duration::from_millis(wait_window_ms);
130    let block_to_completion = params.block_to_completion || params.wait;
131    let timeout = params.timeout;
132    let storage_dir = crate::bash_background::storage_dir(ctx.config().storage_dir.as_deref());
133    let project_root = ctx.config().project_root.clone();
134    let task_id_for_poll = task_id.clone();
135    let request_id_for_poll = request_id.clone();
136    let session_id_for_poll = session_id.clone();
137
138    let mut poll: PendingResponsePoll = Box::new(move |ctx| {
139        let Some(snapshot) = poll_bash_status(
140            ctx,
141            &task_id_for_poll,
142            &session_id_for_poll,
143            project_root.as_deref(),
144            &storage_dir,
145            RUNNING_OUTPUT_PREVIEW_BYTES,
146        ) else {
147            return Some(task_not_found_response(
148                &request_id_for_poll,
149                &task_id_for_poll,
150            ));
151        };
152
153        match decide_bash_step(
154            snapshot,
155            deadline,
156            block_to_completion,
157            Instant::now(),
158            &request_id_for_poll,
159        ) {
160            BashStep::Done(response) => Some(response),
161            BashStep::Promote => Some(promote_bash(
162                ctx,
163                &task_id_for_poll,
164                &session_id_for_poll,
165                project_root.as_deref(),
166                timeout,
167                wait_window_ms,
168                &request_id_for_poll,
169            )),
170            BashStep::Wait => None,
171        }
172    });
173
174    if let Some(response) = poll(ctx) {
175        return DispatchOutcome::Immediate(response);
176    }
177
178    DispatchOutcome::Deferred(PendingResponse {
179        request_id,
180        session_id,
181        attach_command,
182        poll,
183    })
184}
185
186pub(crate) fn poll_bash_status(
187    ctx: &AppContext,
188    task_id: &str,
189    session_id: &str,
190    project_root: Option<&Path>,
191    storage_dir: &Path,
192    preview_bytes: usize,
193) -> Option<BgTaskSnapshot> {
194    ctx.bash_background().status(
195        task_id,
196        session_id,
197        project_root,
198        Some(storage_dir),
199        preview_bytes,
200    )
201}
202
203pub(crate) enum BashStep {
204    Done(Response),
205    Promote,
206    Wait,
207}
208
209pub(crate) fn task_not_found_response(request_id: &str, task_id: &str) -> Response {
210    Response::error(
211        request_id,
212        "task_not_found",
213        format!("background task not found: {task_id}"),
214    )
215}
216
217pub(crate) fn decide_bash_step(
218    snapshot: BgTaskSnapshot,
219    deadline: Instant,
220    block_to_completion: bool,
221    now: Instant,
222    request_id: &str,
223) -> BashStep {
224    if snapshot.info.status.is_terminal() {
225        BashStep::Done(foreground_result_response(request_id, snapshot))
226    } else if !block_to_completion && now >= deadline {
227        BashStep::Promote
228    } else {
229        BashStep::Wait
230    }
231}
232
233pub(crate) fn promote_bash(
234    ctx: &AppContext,
235    task_id: &str,
236    session_id: &str,
237    _project_root: Option<&Path>,
238    timeout: Option<u64>,
239    wait_window_ms: u64,
240    request_id: &str,
241) -> Response {
242    match ctx.bash_background().promote(task_id, session_id) {
243        Ok(_) => promotion_response(request_id, task_id, timeout, wait_window_ms),
244        Err(message) if message.contains("not found") => {
245            Response::error(request_id, "task_not_found", message)
246        }
247        Err(message) => Response::error(request_id, "execution_failed", message),
248    }
249}
250
251fn foreground_result_response(request_id: &str, snapshot: BgTaskSnapshot) -> Response {
252    let output = format_foreground_result(&snapshot);
253    let timed_out = snapshot.info.status == BgTaskStatus::TimedOut;
254    Response::success(
255        request_id,
256        json!({
257            "output": output,
258            "task_id": snapshot.info.task_id,
259            "status": snapshot.info.status,
260            "mode": snapshot.info.mode,
261            "exit_code": snapshot.exit_code,
262            "output_preview": snapshot.output_preview,
263            "output_truncated": snapshot.output_truncated,
264            "truncated": snapshot.output_truncated,
265            "output_path": snapshot.output_path,
266            "timed_out": timed_out,
267            "duration_ms": snapshot.info.duration_ms,
268        }),
269    )
270}
271
272fn background_launch_response(request_id: &str, task_id: &str, is_pty: bool) -> Response {
273    Response::success(
274        request_id,
275        json!({
276            "output": format_background_launch(task_id, is_pty),
277            "task_id": task_id,
278            "status": "running",
279            "mode": if is_pty { "pty" } else { "pipes" },
280        }),
281    )
282}
283
284fn promotion_response(
285    request_id: &str,
286    task_id: &str,
287    timeout: Option<u64>,
288    wait_window_ms: u64,
289) -> Response {
290    Response::success(
291        request_id,
292        json!({
293            "output": format_promotion_message(task_id, timeout, wait_window_ms),
294            "task_id": task_id,
295            "status": "running",
296        }),
297    )
298}
299
300fn parse_params(req: &RawRequest) -> Option<BashOrchestrateParams> {
301    let raw_params = req
302        .params
303        .get("params")
304        .cloned()
305        .unwrap_or_else(|| req.params.clone());
306    serde_json::from_value::<BashOrchestrateParams>(raw_params).ok()
307}
308
309pub(crate) fn resolve_foreground_wait_window_ms(configured: u64) -> u64 {
310    std::env::var(TEST_FOREGROUND_WAIT_ENV)
311        .ok()
312        .and_then(|raw| raw.parse::<u64>().ok())
313        .unwrap_or(configured)
314}
315
316pub(crate) fn select_foreground_wait_window_ms(
317    configured: u64,
318    timeout: Option<u64>,
319    wait: bool,
320) -> u64 {
321    if wait {
322        timeout.unwrap_or(DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS)
323    } else {
324        resolve_foreground_wait_window_ms(configured)
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use crate::bash_background::persistence::BgMode;
332    use crate::bash_background::registry::BgTaskSnapshot;
333    use crate::bash_background::BgTaskInfo;
334
335    fn snapshot(
336        output_preview: &str,
337        output_truncated: bool,
338        output_path: Option<&str>,
339        status: BgTaskStatus,
340        exit_code: Option<i32>,
341    ) -> BgTaskSnapshot {
342        BgTaskSnapshot {
343            info: BgTaskInfo {
344                task_id: "bash-test".to_string(),
345                status,
346                command: "echo test".to_string(),
347                mode: BgMode::Pipes,
348                started_at: 0,
349                duration_ms: Some(1),
350            },
351            exit_code,
352            child_pid: None,
353            workdir: "/tmp".to_string(),
354            output_preview: output_preview.to_string(),
355            output_truncated,
356            output_path: output_path.map(str::to_string),
357            stderr_path: None,
358            pty_rows: None,
359            pty_cols: None,
360            pty_screen: None,
361        }
362    }
363
364    #[test]
365    fn decide_bash_step_returns_done_for_terminal_snapshot_even_at_deadline() {
366        let snapshot = snapshot("done", false, None, BgTaskStatus::Completed, Some(0));
367        let now = Instant::now();
368
369        match decide_bash_step(snapshot, now, false, now, "req-terminal") {
370            BashStep::Done(response) => {
371                assert_eq!(response.id, "req-terminal");
372                assert!(response.success);
373                assert_eq!(response.data["status"], json!("completed"));
374                assert_eq!(response.data["output"], json!("done"));
375            }
376            BashStep::Promote => panic!("terminal snapshot should not promote"),
377            BashStep::Wait => panic!("terminal snapshot should not wait"),
378        }
379    }
380
381    #[test]
382    fn decide_bash_step_promotes_at_deadline_when_not_blocking() {
383        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
384        let now = Instant::now();
385
386        match decide_bash_step(snapshot, now, false, now, "req-promote") {
387            BashStep::Promote => {}
388            BashStep::Done(_) => panic!("running snapshot should not finish"),
389            BashStep::Wait => panic!("deadline should promote when not blocking"),
390        }
391    }
392
393    #[test]
394    fn decide_bash_step_waits_before_deadline() {
395        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
396        let now = Instant::now();
397
398        match decide_bash_step(
399            snapshot,
400            now + Duration::from_millis(1),
401            false,
402            now,
403            "req-wait",
404        ) {
405            BashStep::Wait => {}
406            BashStep::Done(_) => panic!("running snapshot should not finish"),
407            BashStep::Promote => panic!("snapshot should wait before the deadline"),
408        }
409    }
410
411    #[test]
412    fn decide_bash_step_never_promotes_when_blocking_to_completion() {
413        let snapshot = snapshot("running", false, None, BgTaskStatus::Running, None);
414        let now = Instant::now();
415
416        match decide_bash_step(snapshot, now, true, now, "req-block") {
417            BashStep::Wait => {}
418            BashStep::Done(_) => panic!("running snapshot should not finish"),
419            BashStep::Promote => panic!("block_to_completion should suppress promotion"),
420        }
421    }
422
423    #[test]
424    fn select_foreground_wait_window_uses_timeout_budget_for_wait_true() {
425        assert_eq!(
426            select_foreground_wait_window_ms(8_000, Some(250), true),
427            250
428        );
429        assert_eq!(
430            select_foreground_wait_window_ms(8_000, None, true),
431            DEFAULT_FOREGROUND_WAIT_TIMEOUT_MS
432        );
433    }
434
435    #[test]
436    fn foreground_result_format_matches_typescript_order() {
437        let snapshot = snapshot(
438            "hello",
439            true,
440            Some("/tmp/aft-output.txt"),
441            BgTaskStatus::TimedOut,
442            Some(124),
443        );
444
445        assert_eq!(
446            format_foreground_result(&snapshot),
447            "hello\n[output truncated; full output at /tmp/aft-output.txt]\n[command timed out]\n[exit code: 124]"
448        );
449    }
450
451    #[test]
452    fn format_seconds_strips_integer_decimal_and_keeps_tenths() {
453        assert_eq!(format_seconds(8_000), "8s");
454        assert_eq!(format_seconds(5_500), "5.5s");
455        assert_eq!(format_seconds(14_999), "15s");
456    }
457
458    #[test]
459    fn promotion_message_matches_opencode_copy() {
460        assert_eq!(
461            format_promotion_message("bash-123", Some(5_500), 8_000),
462            "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."
463        );
464    }
465
466    #[test]
467    fn background_launch_messages_match_opencode_copy() {
468        assert_eq!(
469            format_background_launch("bash-bg", false),
470            "Background task started: bash-bg. A completion reminder will be delivered automatically; don't poll bash_status."
471        );
472        assert_eq!(
473            format_background_launch("bash-pty", true),
474            "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."
475        );
476    }
477}