Skip to main content

aft/db/
bash_tasks.rs

1use rusqlite::{params, Connection, OptionalExtension, Row};
2
3#[derive(Debug, Clone)]
4pub struct BashTaskRow {
5    pub harness: String,
6    pub session_id: String,
7    pub task_id: String,
8    pub project_key: String,
9    pub command: String,
10    pub cwd: String,
11    pub status: String,
12    pub exit_code: Option<i32>,
13    pub pid: Option<i64>,
14    pub pgid: Option<i64>,
15    pub started_at: i64,
16    pub completed_at: Option<i64>,
17    pub stdout_path: Option<String>,
18    pub stderr_path: Option<String>,
19    pub compressed: bool,
20    pub timeout_ms: Option<i64>,
21    pub completion_delivered: bool,
22    pub output_bytes: Option<i64>,
23    pub metadata: String,
24}
25
26pub fn upsert_bash_task(conn: &Connection, row: &BashTaskRow) -> rusqlite::Result<()> {
27    conn.execute(
28        "INSERT INTO bash_tasks (
29            harness, session_id, task_id, project_key, command, cwd, status,
30            exit_code, pid, pgid, started_at, completed_at, stdout_path, stderr_path,
31            compressed, timeout_ms, completion_delivered, output_bytes, metadata
32         ) VALUES (
33            ?1, ?2, ?3, ?4, ?5, ?6, ?7,
34            ?8, ?9, ?10, ?11, ?12, ?13, ?14,
35            ?15, ?16, ?17, ?18, ?19
36         )
37         ON CONFLICT(harness, session_id, task_id) DO UPDATE SET
38            project_key = excluded.project_key,
39            command = excluded.command,
40            cwd = excluded.cwd,
41            status = excluded.status,
42            exit_code = excluded.exit_code,
43            pid = excluded.pid,
44            pgid = excluded.pgid,
45            started_at = excluded.started_at,
46            completed_at = excluded.completed_at,
47            stdout_path = excluded.stdout_path,
48            stderr_path = excluded.stderr_path,
49            compressed = excluded.compressed,
50            timeout_ms = excluded.timeout_ms,
51            completion_delivered = excluded.completion_delivered,
52            output_bytes = excluded.output_bytes,
53            metadata = excluded.metadata",
54        params![
55            row.harness,
56            row.session_id,
57            row.task_id,
58            row.project_key,
59            row.command,
60            row.cwd,
61            row.status,
62            row.exit_code,
63            row.pid,
64            row.pgid,
65            row.started_at,
66            row.completed_at,
67            row.stdout_path,
68            row.stderr_path,
69            row.compressed,
70            row.timeout_ms,
71            row.completion_delivered,
72            row.output_bytes,
73            row.metadata,
74        ],
75    )?;
76    Ok(())
77}
78
79pub fn delete_delivered_terminal_bash_task(
80    conn: &Connection,
81    harness: &str,
82    session_id: &str,
83    task_id: &str,
84    reason: &str,
85) -> rusqlite::Result<usize> {
86    let deleted = conn.execute(
87        "DELETE FROM bash_tasks
88         WHERE harness = ?1 AND session_id = ?2 AND task_id = ?3
89           AND completion_delivered = 1
90           AND status IN ('completed', 'failed', 'killed', 'timed_out', 'fate_unknown')",
91        params![harness, session_id, task_id],
92    )?;
93    // A row can produce this warning only once: retries affect zero rows after
94    // the first successful DELETE, preventing a cleanup loop from flooding logs.
95    if deleted > 0 {
96        crate::slog_warn!("bash task row deleted: task_id={task_id} reason={reason}");
97    }
98    Ok(deleted)
99}
100
101pub fn delete_bash_task(
102    conn: &Connection,
103    harness: &str,
104    session_id: &str,
105    task_id: &str,
106) -> rusqlite::Result<usize> {
107    conn.execute(
108        "DELETE FROM bash_tasks
109         WHERE harness = ?1 AND session_id = ?2 AND task_id = ?3",
110        params![harness, session_id, task_id],
111    )
112}
113
114pub fn get_bash_task(
115    conn: &Connection,
116    harness: &str,
117    session_id: &str,
118    task_id: &str,
119) -> rusqlite::Result<Option<BashTaskRow>> {
120    conn.query_row(
121        "SELECT harness, session_id, task_id, project_key, command, cwd, status,
122                exit_code, pid, pgid, started_at, completed_at, stdout_path, stderr_path,
123                compressed, timeout_ms, completion_delivered, output_bytes, metadata
124         FROM bash_tasks
125         WHERE harness = ?1 AND session_id = ?2 AND task_id = ?3",
126        params![harness, session_id, task_id],
127        map_bash_task_row,
128    )
129    .optional()
130}
131
132const SESSION_TASKS_SQL: &str =
133    "SELECT harness, session_id, task_id, project_key, command, cwd, status,
134                exit_code, pid, pgid, started_at, completed_at, stdout_path, stderr_path,
135                compressed, timeout_ms, completion_delivered, output_bytes, metadata
136         FROM bash_tasks
137         WHERE harness = ?1 AND session_id = ?2";
138
139pub fn list_bash_tasks_for_session(
140    conn: &Connection,
141    harness: &str,
142    session_id: &str,
143) -> rusqlite::Result<Vec<BashTaskRow>> {
144    let mut stmt = conn.prepare(SESSION_TASKS_SQL)?;
145    let mut rows = stmt
146        .query_map(params![harness, session_id], map_bash_task_row)?
147        .collect::<rusqlite::Result<Vec<_>>>()?;
148    // Task rows carry large command/metadata payloads. Sorting them in SQLite
149    // spills entire rows to a temporary file for long-lived sessions. Keep the
150    // legacy integer/BINARY order, but sort the already-required result vector.
151    rows.sort_by(|a, b| (a.started_at, &a.task_id).cmp(&(b.started_at, &b.task_id)));
152    Ok(rows)
153}
154
155pub fn list_bash_tasks_by_id(
156    conn: &Connection,
157    harness: &str,
158    task_id: &str,
159) -> rusqlite::Result<Vec<BashTaskRow>> {
160    let mut stmt = conn.prepare(
161        "SELECT harness, session_id, task_id, project_key, command, cwd, status,
162                exit_code, pid, pgid, started_at, completed_at, stdout_path, stderr_path,
163                compressed, timeout_ms, completion_delivered, output_bytes, metadata
164         FROM bash_tasks
165         WHERE harness = ?1 AND task_id = ?2
166         ORDER BY started_at DESC",
167    )?;
168    let rows = stmt
169        .query_map(params![harness, task_id], map_bash_task_row)?
170        .collect();
171    rows
172}
173
174pub fn list_replayable_bash_tasks_for_project(
175    conn: &Connection,
176    harness: &str,
177    project_key: &str,
178) -> rusqlite::Result<Vec<BashTaskRow>> {
179    let mut stmt = conn.prepare(
180        "SELECT harness, session_id, task_id, project_key, command, cwd, status,
181                exit_code, pid, pgid, started_at, completed_at, stdout_path, stderr_path,
182                compressed, timeout_ms, completion_delivered, output_bytes, metadata
183         FROM bash_tasks
184         WHERE harness = ?1 AND project_key = ?2
185           AND (status NOT IN ('completed', 'failed', 'killed', 'timed_out', 'fate_unknown')
186                OR completion_delivered = 0)
187         ORDER BY started_at ASC, task_id ASC",
188    )?;
189    let rows = stmt
190        .query_map(params![harness, project_key], map_bash_task_row)?
191        .collect();
192    rows
193}
194
195pub fn find_bash_task_for_project(
196    conn: &Connection,
197    harness: &str,
198    project_key: &str,
199    task_id: &str,
200) -> rusqlite::Result<Option<BashTaskRow>> {
201    conn.query_row(
202        "SELECT harness, session_id, task_id, project_key, command, cwd, status,
203                exit_code, pid, pgid, started_at, completed_at, stdout_path, stderr_path,
204                compressed, timeout_ms, completion_delivered, output_bytes, metadata
205         FROM bash_tasks
206         WHERE harness = ?1 AND project_key = ?2 AND task_id = ?3
207         ORDER BY started_at DESC
208         LIMIT 1",
209        params![harness, project_key, task_id],
210        map_bash_task_row,
211    )
212    .optional()
213}
214
215fn map_bash_task_row(row: &Row<'_>) -> rusqlite::Result<BashTaskRow> {
216    Ok(BashTaskRow {
217        harness: row.get(0)?,
218        session_id: row.get(1)?,
219        task_id: row.get(2)?,
220        project_key: row.get(3)?,
221        command: row.get(4)?,
222        cwd: row.get(5)?,
223        status: row.get(6)?,
224        exit_code: row.get(7)?,
225        pid: row.get(8)?,
226        pgid: row.get(9)?,
227        started_at: row.get(10)?,
228        completed_at: row.get(11)?,
229        stdout_path: row.get(12)?,
230        stderr_path: row.get(13)?,
231        compressed: row.get::<_, i64>(14)? != 0,
232        timeout_ms: row.get(15)?,
233        completion_delivered: row.get::<_, i64>(16)? != 0,
234        output_bytes: row.get(17)?,
235        metadata: row.get::<_, Option<String>>(18)?.unwrap_or_default(),
236    })
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242
243    #[test]
244    fn session_history_preserves_sqlite_order_without_a_temp_sort() {
245        let temp = tempfile::tempdir().unwrap();
246        let conn = crate::db::open(&temp.path().join("aft.db")).unwrap();
247        for (task, started, status) in [
248            ("é", 4, "running"),
249            ("a", 4, "completed"),
250            ("z", -1, "failed"),
251            ("A", 4, "failed"),
252            ("aa", 4, "running"),
253            ("first", i64::MIN, "completed"),
254        ] {
255            conn.execute("INSERT INTO bash_tasks
256                (harness, session_id, task_id, project_key, command, cwd, status, started_at, metadata)
257                VALUES ('opencode', 'session', ?1, 'project', ?2, '.', ?3, ?4, ?5)",
258                params![task, "command".repeat(8192), status, started, format!("metadata-{task}")]).unwrap();
259        }
260        let legacy = conn
261            .prepare(
262                "SELECT harness, session_id, task_id, project_key, command, cwd, status,
263            exit_code, pid, pgid, started_at, completed_at, stdout_path, stderr_path,
264            compressed, timeout_ms, completion_delivered, output_bytes, metadata
265            FROM bash_tasks WHERE harness = ?1 AND session_id = ?2
266            ORDER BY started_at ASC, task_id ASC",
267            )
268            .unwrap()
269            .query_map(params!["opencode", "session"], map_bash_task_row)
270            .unwrap()
271            .collect::<rusqlite::Result<Vec<_>>>()
272            .unwrap();
273        let actual = list_bash_tasks_for_session(&conn, "opencode", "session").unwrap();
274        assert_eq!(format!("{actual:?}"), format!("{legacy:?}"));
275        assert_eq!(
276            actual
277                .iter()
278                .map(|r| r.task_id.as_str())
279                .collect::<Vec<_>>(),
280            ["first", "z", "A", "a", "aa", "é"]
281        );
282        assert!(list_bash_tasks_for_session(&conn, "other", "session")
283            .unwrap()
284            .is_empty());
285        let plan = conn
286            .prepare(&format!("EXPLAIN QUERY PLAN {SESSION_TASKS_SQL}"))
287            .unwrap()
288            .query_map(params!["opencode", "session"], |r| r.get::<_, String>(3))
289            .unwrap()
290            .collect::<rusqlite::Result<Vec<_>>>()
291            .unwrap()
292            .join("\n");
293        assert!(
294            !plan.contains("TEMP B-TREE"),
295            "session history must not spill task rows: {plan}"
296        );
297    }
298}