Skip to main content

aft/bash_background/
persistence.rs

1use std::fs::{self, File, OpenOptions};
2use std::io::{self, Read, Write};
3use std::path::{Path, PathBuf};
4use std::time::{SystemTime, UNIX_EPOCH};
5
6use serde::{Deserialize, Serialize};
7
8use crate::backup::hash_session;
9use crate::bash_permissions::PermissionAsk;
10use crate::db::bash_tasks::BashTaskRow;
11
12use super::BgTaskStatus;
13
14pub const SCHEMA_VERSION: u32 = 5;
15
16#[derive(Debug, Clone)]
17pub struct TaskPaths {
18    pub dir: PathBuf,
19    pub json: PathBuf,
20    pub stdout: PathBuf,
21    pub stderr: PathBuf,
22    pub exit: PathBuf,
23    pub pty: PathBuf,
24    pub sandbox_unavailable: PathBuf,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
28#[serde(rename_all = "lowercase")]
29pub enum BgMode {
30    #[default]
31    Pipes,
32    Pty,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct PersistedTask {
37    pub schema_version: u32,
38    pub task_id: String,
39    pub session_id: String,
40    pub command: String,
41    #[serde(default)]
42    pub mode: BgMode,
43    pub workdir: PathBuf,
44    #[serde(default)]
45    pub project_root: Option<PathBuf>,
46    pub status: BgTaskStatus,
47    pub started_at: u64,
48    pub finished_at: Option<u64>,
49    pub duration_ms: Option<u64>,
50    pub timeout_ms: Option<u64>,
51    pub exit_code: Option<i32>,
52    pub child_pid: Option<u32>,
53    pub pgid: Option<i32>,
54    pub completion_delivered: bool,
55    #[serde(default = "default_notify_on_completion")]
56    pub notify_on_completion: bool,
57    /// Per-call output compression opt-in. Defaults to `true` so existing
58    /// behavior (compression when `experimental.bash.compress=true`) is
59    /// unchanged. Agents can pass `compressed: false` to disable compression
60    /// for a single bash call without flipping the global flag.
61    #[serde(default = "default_compressed")]
62    pub compressed: bool,
63    #[serde(default)]
64    pub pty_rows: Option<u16>,
65    #[serde(default)]
66    pub pty_cols: Option<u16>,
67    /// Dangerous-command findings retained when native containment replaces prompts.
68    #[serde(default, skip_serializing_if = "Vec::is_empty")]
69    pub scanner_report: Vec<PermissionAsk>,
70    /// True when this task was launched through the native sandbox.
71    #[serde(default)]
72    pub sandbox_native: bool,
73    /// Private per-task temp directory created as part of the native profile.
74    #[serde(default, skip_serializing_if = "Option::is_none")]
75    pub sandbox_temp_dir: Option<PathBuf>,
76    pub status_reason: Option<String>,
77}
78
79fn default_notify_on_completion() -> bool {
80    true
81}
82
83fn default_compressed() -> bool {
84    true
85}
86
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub enum ExitMarker {
89    Code(i32),
90    Killed,
91}
92
93impl PersistedTask {
94    pub fn starting(
95        task_id: String,
96        session_id: String,
97        command: String,
98        workdir: PathBuf,
99        project_root: Option<PathBuf>,
100        timeout_ms: Option<u64>,
101        notify_on_completion: bool,
102        compressed: bool,
103    ) -> Self {
104        Self {
105            schema_version: SCHEMA_VERSION,
106            task_id,
107            session_id,
108            command,
109            mode: BgMode::Pipes,
110            workdir,
111            project_root,
112            status: BgTaskStatus::Starting,
113            started_at: unix_millis(),
114            finished_at: None,
115            duration_ms: None,
116            timeout_ms,
117            exit_code: None,
118            child_pid: None,
119            pgid: None,
120            completion_delivered: !notify_on_completion,
121            notify_on_completion,
122            compressed,
123            pty_rows: None,
124            pty_cols: None,
125            scanner_report: Vec::new(),
126            sandbox_native: false,
127            sandbox_temp_dir: None,
128            status_reason: None,
129        }
130    }
131
132    pub fn is_terminal(&self) -> bool {
133        self.status.is_terminal()
134    }
135
136    pub fn mark_running(&mut self, child_pid: u32, pgid: i32) {
137        self.status = BgTaskStatus::Running;
138        self.child_pid = Some(child_pid);
139        self.pgid = Some(pgid);
140    }
141
142    pub fn mark_terminal(
143        &mut self,
144        status: BgTaskStatus,
145        exit_code: Option<i32>,
146        reason: Option<String>,
147    ) {
148        let finished_at = unix_millis();
149        self.status = status;
150        self.exit_code = exit_code;
151        self.finished_at = Some(finished_at);
152        self.duration_ms = Some(finished_at.saturating_sub(self.started_at));
153        self.child_pid = None;
154        self.status_reason = reason;
155        self.completion_delivered = !self.notify_on_completion;
156    }
157
158    pub fn to_bash_task_row(
159        &self,
160        harness: &str,
161        paths: &TaskPaths,
162    ) -> Result<BashTaskRow, serde_json::Error> {
163        let project_root = self.project_root.as_deref().unwrap_or(&self.workdir);
164        let output_bytes = capture_output_bytes(&self.mode, paths);
165        let stdout_path = match self.mode {
166            BgMode::Pipes => Some(paths.stdout.display().to_string()),
167            BgMode::Pty => Some(paths.pty.display().to_string()),
168        };
169        let stderr_path = match self.mode {
170            BgMode::Pipes => Some(paths.stderr.display().to_string()),
171            BgMode::Pty => None,
172        };
173        let mut metadata = self.clone();
174        metadata.schema_version = SCHEMA_VERSION;
175        Ok(BashTaskRow {
176            harness: harness.to_string(),
177            session_id: self.session_id.clone(),
178            task_id: self.task_id.clone(),
179            project_key: crate::path_identity::project_scope_key(project_root),
180            command: self.command.clone(),
181            cwd: self.workdir.display().to_string(),
182            status: status_name(&self.status).to_string(),
183            exit_code: self.exit_code,
184            pid: self.child_pid.map(i64::from),
185            pgid: self.pgid.map(i64::from),
186            started_at: self.started_at as i64,
187            completed_at: self.finished_at.map(|value| value as i64),
188            stdout_path,
189            stderr_path,
190            compressed: self.compressed,
191            timeout_ms: self.timeout_ms.map(|value| value as i64),
192            completion_delivered: self.completion_delivered,
193            output_bytes,
194            metadata: serde_json::to_string(&metadata)?,
195        })
196    }
197}
198
199impl From<BashTaskRow> for PersistedTask {
200    fn from(row: BashTaskRow) -> Self {
201        if let Ok(task) = serde_json::from_str::<PersistedTask>(&row.metadata) {
202            return task;
203        }
204
205        let status = match row.status.as_str() {
206            "starting" => BgTaskStatus::Starting,
207            "running" => BgTaskStatus::Running,
208            "killing" => BgTaskStatus::Killing,
209            "completed" => BgTaskStatus::Completed,
210            "failed" => BgTaskStatus::Failed,
211            "killed" => BgTaskStatus::Killed,
212            "timed_out" => BgTaskStatus::TimedOut,
213            _ => BgTaskStatus::Failed,
214        };
215        let started_at = u64::try_from(row.started_at).unwrap_or_default();
216        let finished_at = row.completed_at.and_then(|value| u64::try_from(value).ok());
217
218        PersistedTask {
219            schema_version: SCHEMA_VERSION,
220            task_id: row.task_id,
221            session_id: row.session_id,
222            command: row.command,
223            mode: BgMode::Pipes,
224            workdir: PathBuf::from(row.cwd),
225            project_root: None,
226            status,
227            started_at,
228            finished_at,
229            duration_ms: finished_at.map(|finished_at| finished_at.saturating_sub(started_at)),
230            timeout_ms: row.timeout_ms.and_then(|value| u64::try_from(value).ok()),
231            exit_code: row.exit_code,
232            child_pid: row.pid.and_then(|value| u32::try_from(value).ok()),
233            pgid: row.pgid.and_then(|value| i32::try_from(value).ok()),
234            completion_delivered: row.completion_delivered,
235            notify_on_completion: !row.completion_delivered,
236            compressed: row.compressed,
237            pty_rows: None,
238            pty_cols: None,
239            scanner_report: Vec::new(),
240            sandbox_native: false,
241            sandbox_temp_dir: None,
242            status_reason: None,
243        }
244    }
245}
246
247fn status_name(status: &BgTaskStatus) -> &'static str {
248    match status {
249        BgTaskStatus::Starting => "starting",
250        BgTaskStatus::Running => "running",
251        BgTaskStatus::Killing => "killing",
252        BgTaskStatus::Completed => "completed",
253        BgTaskStatus::Failed => "failed",
254        BgTaskStatus::Killed => "killed",
255        BgTaskStatus::TimedOut => "timed_out",
256    }
257}
258
259fn capture_output_bytes(mode: &BgMode, paths: &TaskPaths) -> Option<i64> {
260    match mode {
261        BgMode::Pipes => {
262            let stdout = fs::metadata(&paths.stdout)
263                .ok()
264                .map(|metadata| metadata.len());
265            let stderr = fs::metadata(&paths.stderr)
266                .ok()
267                .map(|metadata| metadata.len());
268            match (stdout, stderr) {
269                (Some(stdout), Some(stderr)) => Some(stdout.saturating_add(stderr) as i64),
270                (Some(bytes), None) | (None, Some(bytes)) => Some(bytes as i64),
271                (None, None) => None,
272            }
273        }
274        BgMode::Pty => fs::metadata(&paths.pty)
275            .ok()
276            .map(|metadata| metadata.len() as i64),
277    }
278}
279
280pub fn session_tasks_dir(storage_dir: &Path, session_id: &str) -> PathBuf {
281    let session_hash = hash_session(session_id);
282    let direct = storage_dir.join("bash-tasks").join(&session_hash);
283    if direct.exists() {
284        return direct;
285    }
286
287    let mut harness_matches = ["opencode", "pi"]
288        .into_iter()
289        .map(|harness| {
290            storage_dir
291                .join(harness)
292                .join("bash-tasks")
293                .join(&session_hash)
294        })
295        .filter(|path| path.exists())
296        .collect::<Vec<_>>();
297    if harness_matches.len() == 1 {
298        return harness_matches.remove(0);
299    }
300
301    direct
302}
303
304pub fn task_paths(storage_dir: &Path, session_id: &str, task_id: &str) -> TaskPaths {
305    let dir = session_tasks_dir(storage_dir, session_id);
306    TaskPaths {
307        json: dir.join(format!("{task_id}.json")),
308        stdout: dir.join(format!("{task_id}.stdout")),
309        stderr: dir.join(format!("{task_id}.stderr")),
310        exit: dir.join(format!("{task_id}.exit")),
311        pty: dir.join(format!("{task_id}.pty")),
312        sandbox_unavailable: dir.join(format!("{task_id}.sandbox-unavailable")),
313        dir,
314    }
315}
316
317pub fn read_task(path: &Path) -> io::Result<PersistedTask> {
318    let content = fs::read_to_string(path)?;
319    let task: PersistedTask = serde_json::from_str(&content).map_err(io::Error::other)?;
320    if !matches!(task.schema_version, 2 | 3 | 4 | SCHEMA_VERSION) {
321        return Err(io::Error::new(
322            io::ErrorKind::InvalidData,
323            format!(
324                "unsupported background task schema_version {} (expected 2, 3, 4, or {SCHEMA_VERSION})",
325                task.schema_version
326            ),
327        ));
328    }
329    Ok(task)
330}
331
332pub fn write_task(path: &Path, task: &PersistedTask) -> io::Result<()> {
333    if let Some(parent) = path.parent() {
334        fs::create_dir_all(parent)?;
335    }
336    let mut upgraded = task.clone();
337    upgraded.schema_version = SCHEMA_VERSION;
338    let content = serde_json::to_vec_pretty(&upgraded).map_err(io::Error::other)?;
339    atomic_write(path, &content, &upgraded.task_id)
340}
341
342pub(super) fn delete_task_bundle(paths: &TaskPaths) -> io::Result<()> {
343    let managed_temp_dir = read_task(&paths.json)
344        .ok()
345        .and_then(|task| task.sandbox_temp_dir)
346        .filter(|path| {
347            crate::sandbox_spawn::is_managed_task_temp_dir(path)
348                && path.parent().is_some_and(|parent| {
349                    let canonical_dir = paths
350                        .dir
351                        .canonicalize()
352                        .unwrap_or_else(|_| paths.dir.clone());
353                    parent == canonical_dir
354                })
355        });
356
357    let mut first_error = None;
358    if let Some(temp_dir) = managed_temp_dir {
359        if let Err(error) = fs::remove_dir_all(temp_dir) {
360            if error.kind() != io::ErrorKind::NotFound {
361                first_error = Some(error);
362            }
363        }
364    }
365    for path in task_bundle_files(paths) {
366        if let Err(error) = remove_file_if_present(&path) {
367            if first_error.is_none() {
368                first_error = Some(error);
369            }
370        }
371    }
372
373    if let Some(error) = first_error {
374        return Err(error);
375    }
376
377    match fs::remove_dir(&paths.dir) {
378        Ok(()) => Ok(()),
379        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
380        Err(error) if error.kind() == io::ErrorKind::DirectoryNotEmpty => Ok(()),
381        Err(error) => Err(error),
382    }
383}
384
385pub fn task_bundle_files(paths: &TaskPaths) -> Vec<PathBuf> {
386    let mut files = vec![
387        paths.json.clone(),
388        paths.stdout.clone(),
389        paths.stderr.clone(),
390        paths.exit.clone(),
391        paths.pty.clone(),
392        paths.sandbox_unavailable.clone(),
393    ];
394    if let Some(stem) = paths.json.file_stem().and_then(|stem| stem.to_str()) {
395        // Background bash writes per-task wrapper/command scripts next to the
396        // capture files as `<task-id>.ps1`, `<task-id>.bat`, or `<task-id>.sh`
397        // depending on platform and shell selection.
398        for extension in ["ps1", "bat", "sh"] {
399            files.push(paths.dir.join(format!("{stem}.{extension}")));
400        }
401        files.push(paths.dir.join(format!("{stem}.sandbox-profile.json")));
402    }
403    files
404}
405
406fn remove_file_if_present(path: &Path) -> io::Result<()> {
407    match fs::remove_file(path) {
408        Ok(()) => Ok(()),
409        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
410        Err(error) => Err(error),
411    }
412}
413
414pub fn update_task<F>(path: &Path, update: F) -> io::Result<PersistedTask>
415where
416    F: FnOnce(&mut PersistedTask),
417{
418    let mut task = read_task(path)?;
419    let original_terminal = task.is_terminal();
420    let original = task.clone();
421    update(&mut task);
422    task.schema_version = SCHEMA_VERSION;
423    if original_terminal {
424        let completion_delivered = task.completion_delivered;
425        task = original;
426        task.completion_delivered = completion_delivered;
427        task.schema_version = SCHEMA_VERSION;
428    }
429    write_task(path, &task)?;
430    Ok(task)
431}
432
433pub fn write_kill_marker_if_absent(path: &Path) -> io::Result<()> {
434    if path.exists() {
435        return Ok(());
436    }
437    atomic_write(path, b"killed", "kill")
438}
439
440pub fn read_exit_marker(path: &Path) -> io::Result<Option<ExitMarker>> {
441    let mut file = match File::open(path) {
442        Ok(file) => file,
443        Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(None),
444        Err(error) => return Err(error),
445    };
446    let mut content = String::new();
447    file.read_to_string(&mut content)?;
448    let content = content.trim();
449    if content.is_empty() {
450        return Ok(None);
451    }
452    if content == "killed" {
453        return Ok(Some(ExitMarker::Killed));
454    }
455    match content.parse::<i32>() {
456        Ok(code) => Ok(Some(ExitMarker::Code(code))),
457        Err(_) => Ok(None),
458    }
459}
460
461pub fn atomic_write(path: &Path, content: &[u8], task_id: &str) -> io::Result<()> {
462    let parent = path.parent().unwrap_or_else(|| Path::new("."));
463    fs::create_dir_all(parent)?;
464    let file_name = path
465        .file_name()
466        .and_then(|name| name.to_str())
467        .unwrap_or("task");
468    let tmp = parent.join(format!(
469        ".{file_name}.tmp.{}.{}",
470        std::process::id(),
471        sanitize_task_id(task_id)
472    ));
473    {
474        let mut file = OpenOptions::new()
475            .create(true)
476            .truncate(true)
477            .write(true)
478            .open(&tmp)?;
479        file.write_all(content)?;
480        file.sync_all()?;
481    }
482    fs::rename(&tmp, path)?;
483    Ok(())
484}
485
486fn sanitize_task_id(task_id: &str) -> String {
487    task_id
488        .chars()
489        .map(|ch| match ch {
490            'a'..='z' | 'A'..='Z' | '0'..='9' | '-' | '_' => ch,
491            _ => '_',
492        })
493        .collect()
494}
495
496pub fn create_capture_file(path: &Path) -> io::Result<File> {
497    if let Some(parent) = path.parent() {
498        fs::create_dir_all(parent)?;
499    }
500    File::create(path)
501}
502
503pub fn unix_millis() -> u64 {
504    SystemTime::now()
505        .duration_since(UNIX_EPOCH)
506        .map(|duration| duration.as_millis() as u64)
507        .unwrap_or(0)
508}
509
510#[cfg(test)]
511mod tests {
512    use std::thread;
513
514    use super::*;
515
516    #[test]
517    fn atomic_write_temp_names_include_task_id() {
518        let dir = tempfile::tempdir().expect("create temp dir");
519        let path = dir.path().join("task.json");
520
521        let left_path = path.clone();
522        let left = thread::spawn(move || atomic_write(&left_path, b"left", "task-left"));
523        let right_path = path.clone();
524        let right = thread::spawn(move || atomic_write(&right_path, b"right", "task-right"));
525
526        left.join().expect("join left").expect("write left");
527        right.join().expect("join right").expect("write right");
528
529        let content = fs::read_to_string(&path).expect("read final content");
530        assert!(content == "left" || content == "right");
531        assert!(!dir
532            .path()
533            .join(format!(".task.json.tmp.{}", std::process::id()))
534            .exists());
535    }
536}