Skip to main content

recursive/tools/
run_background.rs

1//! `run_background` and `check_background`: spawn long-running commands
2//! asynchronously and poll their status later.
3//!
4//! `run_background` returns a job ID immediately so the agent can continue
5//! working while the command runs. `check_background` retrieves the output
6//! (or partial output) of a previously spawned job.
7//!
8//! Jobs are stored in a shared `BackgroundJobManager` and cleaned up after
9//! a configurable TTL or when explicitly checked.
10
11use async_trait::async_trait;
12use serde_json::{json, Value};
13use std::collections::HashMap;
14use std::path::PathBuf;
15use std::process::Stdio;
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18use tokio::io::AsyncReadExt;
19use tokio::process::Command;
20use tokio::sync::Mutex;
21use tokio::sync::Notify;
22use tokio::time::timeout;
23
24use super::resolve_within;
25use super::Tool;
26use crate::error::{Error, Result};
27use crate::llm::ToolSpec;
28
29/// Maximum bytes of stdout/stderr to capture per job.
30const MAX_OUTPUT_BYTES: usize = 128 * 1024;
31
32/// Maximum bytes of a watched file to surface in a single event-watch wake.
33/// Keeps the injected prompt bounded; remaining bytes wake on the next poll.
34pub(crate) const WATCH_CHUNK_BYTES: usize = 16 * 1024;
35
36/// Default timeout for a background job (if none specified).
37const DEFAULT_JOB_TIMEOUT: u64 = 3600; // 1 hour
38
39/// How long a completed/failed job's output is kept before cleanup.
40const JOB_TTL: Duration = Duration::from_secs(3600);
41
42/// State of a background job.
43#[derive(Debug, Clone)]
44pub enum JobState {
45    Running,
46    Completed {
47        stdout: String,
48        stderr: String,
49        exit_code: i32,
50    },
51    Failed {
52        message: String,
53    },
54    TimedOut,
55}
56
57impl JobState {
58    /// True if the job will never change state again.
59    pub fn is_terminal(&self) -> bool {
60        matches!(
61            self,
62            JobState::Completed { .. } | JobState::Failed { .. } | JobState::TimedOut
63        )
64    }
65}
66
67/// A single background job.
68pub struct Job {
69    pub state: JobState,
70    pub created_at: Instant,
71}
72
73/// A file registered for mid-run event wakes. The TUI loop arbiter polls the
74/// file and wakes the agent when bytes are appended past `offset`.
75#[derive(Debug, Clone)]
76pub struct WatchTarget {
77    pub path: PathBuf,
78    pub offset: u64,
79}
80
81/// Agent-requested loop control written by the `stop_loop` tool and read by
82/// the TUI loop arbiter between turns. Lets the agent end the event-driven
83/// loop itself (so the user can say "stop" in natural language and the agent
84/// calls `stop_loop`), making the loop lifecycle transparent to the user.
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum LoopControl {
87    /// Stop the loop after the current turn and return to interactive mode.
88    Stop,
89}
90
91/// Shared manager for background jobs.
92pub struct BackgroundJobManager {
93    jobs: HashMap<String, Job>,
94    next_id: u64,
95    /// Notifies waiters when a job enters a terminal state
96    /// (Completed, Failed, TimedOut). Used by the TUI loop arbiter
97    /// to wake on background job completion.
98    completed_notify: Arc<Notify>,
99    /// Optional file registered via the `watch_file` tool for mid-run
100    /// event wakes. The loop arbiter polls it and wakes the agent when
101    /// new bytes appear past `offset`.
102    watch: Option<WatchTarget>,
103    /// Agent-requested loop control (set by the `stop_loop` tool). The loop
104    /// arbiter drains this between turns so the agent can end the loop itself.
105    loop_control: Option<LoopControl>,
106}
107
108impl Default for BackgroundJobManager {
109    fn default() -> Self {
110        Self::new()
111    }
112}
113
114impl BackgroundJobManager {
115    pub fn new() -> Self {
116        Self {
117            jobs: HashMap::new(),
118            next_id: 1,
119            completed_notify: Arc::new(Notify::new()),
120            watch: None,
121            loop_control: None,
122        }
123    }
124
125    /// Return a clone of the completed-notify handle.
126    ///
127    /// Used by the TUI loop arbiter to `select!` on job completion.
128    /// When a background job enters a terminal state (Completed, Failed,
129    /// TimedOut), `notify_one()` is called. `notify_one()` (rather than
130    /// `notify_waiters()`) stores a permit so a completion that happens
131    /// while the arbiter is mid-turn (not yet polling) is not lost — the
132    /// next `notified()` poll consumes the stored permit.
133    pub fn completed_notify(&self) -> Arc<Notify> {
134        self.completed_notify.clone()
135    }
136
137    /// Generate a new unique job ID.
138    fn next_id(&mut self) -> String {
139        let id = self.next_id;
140        self.next_id += 1;
141        format!("bg-{id}")
142    }
143
144    /// Insert a new job and return its ID.
145    pub fn insert(&mut self, job: Job) -> String {
146        let id = self.next_id();
147        self.jobs.insert(id.clone(), job);
148        id
149    }
150
151    /// Get a reference to a job's state.
152    pub fn get_state(&self, id: &str) -> Option<JobState> {
153        self.jobs.get(id).map(|j| j.state.clone())
154    }
155
156    /// Update a job's state.
157    fn update(&mut self, id: &str, state: JobState) {
158        let is_terminal = state.is_terminal();
159        if let Some(job) = self.jobs.get_mut(id) {
160            job.state = state;
161        }
162        if is_terminal {
163            self.completed_notify.notify_one();
164        }
165    }
166
167    /// Remove stale jobs (older than TTL).
168    fn cleanup(&mut self) {
169        // checked_sub: when the elapsed time since the Instant epoch is less than
170        // JOB_TTL (e.g. shortly after system boot), `now - TTL` would underflow and
171        // panic. In that case no job can possibly be older than the TTL, so keep them
172        // all (cutoff = now → `created_at > now` is false for every job, but that's
173        // fine — nothing is evicted, which is correct since nothing is stale).
174        let cutoff = Instant::now()
175            .checked_sub(JOB_TTL)
176            .unwrap_or_else(Instant::now);
177        self.jobs.retain(|_, job| job.created_at > cutoff);
178    }
179
180    /// Remove all jobs immediately.
181    pub fn clear(&mut self) {
182        self.jobs.clear();
183    }
184
185    /// Register a file for mid-run event wakes. `offset` is the byte offset to
186    /// start reading from (typically the file size at registration time, so
187    /// only bytes appended afterwards wake the agent). Replaces any prior
188    /// watch.
189    pub fn set_watch(&mut self, path: PathBuf, offset: u64) {
190        self.watch = Some(WatchTarget { path, offset });
191    }
192
193    /// Clear any registered file watch.
194    pub fn clear_watch(&mut self) {
195        self.watch = None;
196    }
197
198    /// Set an agent-requested loop control (e.g. `Stop`). Replaces any prior
199    /// pending control. Read and consumed by the loop arbiter between turns.
200    pub fn set_loop_control(&mut self, ctrl: LoopControl) {
201        self.loop_control = Some(ctrl);
202    }
203
204    /// Take (consume) any pending loop-control request. Returns `None` when
205    /// nothing is pending.
206    pub fn take_loop_control(&mut self) -> Option<LoopControl> {
207        self.loop_control.take()
208    }
209
210    /// Poll the registered watch file for new bytes past the stored offset.
211    /// On success, advances the offset and returns the new content (capped at
212    /// `WATCH_CHUNK_BYTES` per call; remaining bytes surface on the next
213    /// poll). Returns `None` when no watch is set, the file is gone, or no new
214    /// bytes are available. Handles log rotation (file shrank) by resetting to
215    /// the start.
216    pub fn poll_watch(&mut self) -> Option<String> {
217        use std::io::{Read, Seek, SeekFrom};
218        let watch = self.watch.as_ref()?;
219        let path = watch.path.clone();
220        let metadata = std::fs::metadata(&path).ok()?;
221        let len = metadata.len();
222        let mut offset = watch.offset;
223        // Log rotation / truncation: reset to start.
224        if len < offset {
225            offset = 0;
226        }
227        if len <= offset {
228            return None;
229        }
230        let mut file = std::fs::File::open(&path).ok()?;
231        file.seek(SeekFrom::Start(offset)).ok()?;
232        let to_read = ((len - offset) as usize).min(WATCH_CHUNK_BYTES);
233        let mut buf = vec![0u8; to_read];
234        let n = file.read(&mut buf).ok()?;
235        if n == 0 {
236            return None;
237        }
238        let new_offset = offset + n as u64;
239        let text = String::from_utf8_lossy(&buf[..n]).into_owned();
240        if let Some(w) = self.watch.as_mut() {
241            w.offset = new_offset;
242        }
243        if text.is_empty() {
244            None
245        } else {
246            Some(text)
247        }
248    }
249
250    /// Remove and return the first completed job, if any.
251    ///
252    /// Returns `Some((job_id, output_string))` where `output_string` includes
253    /// the exit code and captured stdout/stderr. Completed jobs are removed
254    /// from the tracked set.
255    ///
256    /// If a terminal job remains after this removal, the completion notify is
257    /// re-armed. `update()` uses `notify_one()`, which stores a single permit:
258    /// when ≥2 jobs reach a terminal state before a waiter polls, the second
259    /// `notify_one()` is a no-op, so a waiter that consumes the one permit
260    /// would never wake for the remaining job. Re-arming here guarantees every
261    /// completed job has a scheduled wake until it is drained (goal-379).
262    pub fn take_completed(&mut self) -> Option<(String, String)> {
263        let id = self.jobs.iter().find_map(|(id, job)| match &job.state {
264            JobState::Completed { .. } | JobState::Failed { .. } | JobState::TimedOut => {
265                Some(id.clone())
266            }
267            JobState::Running => None,
268        })?;
269        let job = self.jobs.remove(&id)?;
270        let output = match &job.state {
271            JobState::Completed {
272                stdout,
273                stderr,
274                exit_code,
275            } => {
276                format!("exit code: {exit_code}\nstdout:\n{stdout}\nstderr:\n{stderr}")
277            }
278            JobState::Failed { message } => {
279                format!("failed: {message}")
280            }
281            JobState::TimedOut => "timed out".to_string(),
282            JobState::Running => unreachable!(), // filtered above
283        };
284        if self.jobs.values().any(|j| j.state.is_terminal()) {
285            self.completed_notify.notify_one();
286        }
287        Some((id, output))
288    }
289}
290
291/// The `run_background` tool: spawn a command and return immediately.
292pub struct RunBackground {
293    root: PathBuf,
294    manager: Arc<Mutex<BackgroundJobManager>>,
295}
296
297impl RunBackground {
298    pub fn new(root: impl Into<PathBuf>, manager: Arc<Mutex<BackgroundJobManager>>) -> Self {
299        Self {
300            root: root.into(),
301            manager,
302        }
303    }
304}
305
306#[async_trait]
307impl Tool for RunBackground {
308    fn spec(&self) -> ToolSpec {
309        ToolSpec {
310            name: "run_background".into(),
311            description:
312                "Run a shell command in the background and return immediately with a job ID. \
313                 Use `check_background` later to retrieve the output. Useful for long-running \
314                 commands (e.g. builds, tests, downloads) that you don't want to block on."
315                    .into(),
316            parameters: json!({
317                "type": "object",
318                "properties": {
319                    "command": {
320                        "type": "string",
321                        "description": "Command line to execute via sh -c"
322                    },
323                    "cwd": {
324                        "type": "string",
325                        "description": "Optional subdirectory (relative to workspace root) to run the command in. Must stay inside the workspace."
326                    },
327                    "env": {
328                        "type": "object",
329                        "description": "Optional extra env vars set for this command only. Values must be strings; non-string values are rejected.",
330                        "additionalProperties": {
331                            "type": "string"
332                        }
333                    },
334                    "timeout_secs": {
335                        "type": "integer",
336                        "description": "Optional timeout in seconds (default 3600, max 86400)",
337                        "default": 3600
338                    }
339                },
340                "required": ["command"]
341            }),
342        }
343    }
344
345    fn is_deferred(&self) -> bool {
346        true
347    }
348
349    async fn execute(&self, args: Value) -> Result<String> {
350        let command = args["command"].as_str().ok_or_else(|| Error::BadToolArgs {
351            name: "run_background".into(),
352            message: "missing `command`".into(),
353        })?;
354
355        let cwd = if let Some(rel) = args.get("cwd").and_then(|v| v.as_str()) {
356            resolve_within(&self.root, rel).map_err(|e| Error::BadToolArgs {
357                name: "run_background".into(),
358                message: format!("cwd: {e}"),
359            })?
360        } else {
361            self.root.clone()
362        };
363
364        let timeout_secs = args
365            .get("timeout_secs")
366            .and_then(|v| v.as_i64())
367            .unwrap_or(DEFAULT_JOB_TIMEOUT as i64)
368            .clamp(1, 86400) as u64;
369
370        let mut cmd = Command::new("/bin/sh");
371        cmd.arg("-c").arg(command);
372        cmd.current_dir(&cwd);
373        cmd.stdout(Stdio::piped());
374        cmd.stderr(Stdio::piped());
375
376        // Apply optional env overrides
377        if let Some(env_map) = args.get("env").and_then(|v| v.as_object()) {
378            for (key, val) in env_map {
379                let val_str = val.as_str().ok_or_else(|| Error::BadToolArgs {
380                    name: "run_background".to_string(),
381                    message: format!("env value for `{key}` must be a string, got {:?}", val),
382                })?;
383                cmd.env(key, val_str);
384            }
385        }
386
387        // Spawn the process
388        let mut child = cmd.spawn().map_err(|e| Error::Tool {
389            name: "run_background".into(),
390            call_id: None,
391            message: format!("spawn failed: {e}"),
392        })?;
393
394        let stdout_handle = child.stdout.take();
395        let stderr_handle = child.stderr.take();
396
397        // Generate a job ID and store it
398        let mut manager = self.manager.lock().await;
399        manager.cleanup();
400        let job_id = manager.insert(Job {
401            state: JobState::Running,
402            created_at: Instant::now(),
403        });
404        drop(manager);
405
406        // Spawn a background task to wait for the process and capture output
407        let manager_clone = self.manager.clone();
408        let job_id_clone = job_id.clone();
409        tokio::spawn(async move {
410            let result = run_background_job(
411                child,
412                stdout_handle,
413                stderr_handle,
414                Duration::from_secs(timeout_secs),
415            )
416            .await;
417
418            let mut mgr = manager_clone.lock().await;
419            mgr.update(&job_id_clone, result);
420        });
421
422        Ok(json!({
423            "job_id": job_id,
424            "status": "spawned",
425            "message": format!("Background job `{}` spawned. Use `check_background` with job_id `{}` to retrieve output.", command, job_id)
426        })
427        .to_string())
428    }
429}
430
431/// The `check_background` tool: poll a previously spawned background job.
432pub struct CheckBackground {
433    manager: Arc<Mutex<BackgroundJobManager>>,
434}
435
436impl CheckBackground {
437    pub fn new(manager: Arc<Mutex<BackgroundJobManager>>) -> Self {
438        Self { manager }
439    }
440}
441
442#[async_trait]
443impl Tool for CheckBackground {
444    fn spec(&self) -> ToolSpec {
445        ToolSpec {
446            name: "check_background".into(),
447            description: "Check the status and output of a previously spawned background job. \
448                 Returns the current state (running, completed, failed, or timed out) \
449                 along with any captured stdout/stderr."
450                .into(),
451            parameters: json!({
452                "type": "object",
453                "properties": {
454                    "job_id": {
455                        "type": "string",
456                        "description": "The job ID returned by `run_background`"
457                    }
458                },
459                "required": ["job_id"]
460            }),
461        }
462    }
463
464    fn is_deferred(&self) -> bool {
465        true
466    }
467
468    async fn execute(&self, args: Value) -> Result<String> {
469        let job_id = args["job_id"].as_str().ok_or_else(|| Error::BadToolArgs {
470            name: "check_background".into(),
471            message: "missing `job_id`".into(),
472        })?;
473
474        let mut manager = self.manager.lock().await;
475        manager.cleanup();
476
477        match manager.get_state(job_id) {
478            None => Ok(json!({
479                "job_id": job_id,
480                "status": "unknown",
481                "message": format!("No job found with ID `{}`. It may have expired or the ID is invalid.", job_id)
482            })
483            .to_string()),
484            Some(JobState::Running) => Ok(json!({
485                "job_id": job_id,
486                "status": "running",
487                "message": "Job is still running. Check again later."
488            })
489            .to_string()),
490            Some(JobState::Completed { stdout, stderr, exit_code }) => {
491                // Clean up completed job after returning its output
492                manager.update(job_id, JobState::Completed {
493                    stdout: String::new(),
494                    stderr: String::new(),
495                    exit_code,
496                });
497                Ok(json!({
498                    "job_id": job_id,
499                    "status": "completed",
500                    "exit_code": exit_code,
501                    "stdout": stdout,
502                    "stderr": stderr
503                })
504                .to_string())
505            }
506            Some(JobState::Failed { message }) => {
507                Ok(json!({
508                    "job_id": job_id,
509                    "status": "failed",
510                    "message": message
511                })
512                .to_string())
513            }
514            Some(JobState::TimedOut) => Ok(json!({
515                "job_id": job_id,
516                "status": "timed_out",
517                "message": "Job exceeded its timeout."
518            })
519            .to_string()),
520        }
521    }
522}
523
524/// Wait for a background process to finish, capturing its output.
525async fn run_background_job(
526    mut child: tokio::process::Child,
527    stdout_opt: Option<tokio::process::ChildStdout>,
528    stderr_opt: Option<tokio::process::ChildStderr>,
529    job_timeout: Duration,
530) -> JobState {
531    let stdout_task = async {
532        let mut out = String::new();
533        if let Some(mut reader) = stdout_opt {
534            let mut buf = [0u8; 8192];
535            let mut total = 0usize;
536            loop {
537                match reader.read(&mut buf).await {
538                    Ok(0) => break,
539                    Ok(n) => {
540                        if total + n > MAX_OUTPUT_BYTES {
541                            let take = MAX_OUTPUT_BYTES.saturating_sub(total);
542                            out.push_str(&String::from_utf8_lossy(&buf[..take]));
543                            out.push_str("\n... [stdout truncated]");
544                            // Drain remaining
545                            let _ = tokio::io::copy(&mut reader, &mut tokio::io::sink()).await;
546                            break;
547                        }
548                        out.push_str(&String::from_utf8_lossy(&buf[..n]));
549                        total += n;
550                    }
551                    Err(_) => break,
552                }
553            }
554        }
555        out
556    };
557
558    let stderr_task = async {
559        let mut err = String::new();
560        if let Some(mut reader) = stderr_opt {
561            let mut buf = [0u8; 8192];
562            let mut total = 0usize;
563            loop {
564                match reader.read(&mut buf).await {
565                    Ok(0) => break,
566                    Ok(n) => {
567                        if total + n > MAX_OUTPUT_BYTES {
568                            let take = MAX_OUTPUT_BYTES.saturating_sub(total);
569                            err.push_str(&String::from_utf8_lossy(&buf[..take]));
570                            err.push_str("\n... [stderr truncated]");
571                            let _ = tokio::io::copy(&mut reader, &mut tokio::io::sink()).await;
572                            break;
573                        }
574                        err.push_str(&String::from_utf8_lossy(&buf[..n]));
575                        total += n;
576                    }
577                    Err(_) => break,
578                }
579            }
580        }
581        err
582    };
583
584    // Wait for the process with a timeout
585    let wait_result = timeout(job_timeout, child.wait()).await;
586
587    match wait_result {
588        Err(_) => {
589            // Timed out — kill the process
590            let _ = child.start_kill();
591            JobState::TimedOut
592        }
593        Ok(Err(e)) => JobState::Failed {
594            message: format!("wait failed: {e}"),
595        },
596        Ok(Ok(status)) => {
597            let stdout = stdout_task.await;
598            let stderr = stderr_task.await;
599            let exit_code = status.code().unwrap_or(-1);
600            JobState::Completed {
601                stdout,
602                stderr,
603                exit_code,
604            }
605        }
606    }
607}
608
609#[cfg(test)]
610#[cfg(not(target_os = "windows"))]
611mod tests {
612    use super::*;
613    use serde_json::json;
614    use tempfile::TempDir;
615    use tokio::time::sleep;
616
617    /// Poll `check_background` until the job is no longer running, with a timeout.
618    async fn poll_until_done(
619        check_tool: &CheckBackground,
620        job_id: &str,
621        max_wait: Duration,
622    ) -> Value {
623        let start = Instant::now();
624        loop {
625            let result = check_tool.execute(json!({"job_id": job_id})).await.unwrap();
626            let parsed: Value = serde_json::from_str(&result).unwrap();
627            if parsed["status"] != "running" {
628                return parsed;
629            }
630            if start.elapsed() > max_wait {
631                panic!("timed out waiting for job {job_id} to complete");
632            }
633            sleep(Duration::from_millis(50)).await;
634        }
635    }
636
637    #[tokio::test]
638    async fn run_and_check_background() {
639        let tmp = TempDir::new().unwrap();
640        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
641
642        let run_tool = RunBackground::new(tmp.path(), manager.clone());
643        let check_tool = CheckBackground::new(manager.clone());
644
645        // Run a quick command
646        let result = run_tool
647            .execute(json!({"command": "echo hello world"}))
648            .await
649            .unwrap();
650
651        let parsed: Value = serde_json::from_str(&result).unwrap();
652        let job_id = parsed["job_id"].as_str().unwrap().to_string();
653        assert_eq!(parsed["status"], "spawned");
654
655        let parsed = poll_until_done(&check_tool, &job_id, Duration::from_secs(5)).await;
656        assert_eq!(parsed["status"], "completed");
657        assert_eq!(parsed["exit_code"], 0);
658        assert!(parsed["stdout"].as_str().unwrap().contains("hello world"));
659    }
660
661    #[tokio::test]
662    async fn background_job_timeout() {
663        let tmp = TempDir::new().unwrap();
664        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
665
666        let run_tool = RunBackground::new(tmp.path(), manager.clone());
667        let check_tool = CheckBackground::new(manager.clone());
668
669        // Run a command that sleeps longer than the timeout
670        let result = run_tool
671            .execute(json!({"command": "sleep 10", "timeout_secs": 1}))
672            .await
673            .unwrap();
674
675        let parsed: Value = serde_json::from_str(&result).unwrap();
676        let job_id = parsed["job_id"].as_str().unwrap().to_string();
677
678        let parsed = poll_until_done(&check_tool, &job_id, Duration::from_secs(5)).await;
679        assert_eq!(parsed["status"], "timed_out");
680    }
681
682    #[tokio::test]
683    async fn background_job_nonzero_exit() {
684        let tmp = TempDir::new().unwrap();
685        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
686
687        let run_tool = RunBackground::new(tmp.path(), manager.clone());
688        let check_tool = CheckBackground::new(manager.clone());
689
690        let result = run_tool
691            .execute(json!({"command": "exit 42"}))
692            .await
693            .unwrap();
694
695        let parsed: Value = serde_json::from_str(&result).unwrap();
696        let job_id = parsed["job_id"].as_str().unwrap().to_string();
697
698        let parsed = poll_until_done(&check_tool, &job_id, Duration::from_secs(5)).await;
699        assert_eq!(parsed["status"], "completed");
700        assert_eq!(parsed["exit_code"], 42);
701    }
702
703    #[tokio::test]
704    async fn check_unknown_job() {
705        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
706        let check_tool = CheckBackground::new(manager);
707
708        let result = check_tool
709            .execute(json!({"job_id": "bg-999"}))
710            .await
711            .unwrap();
712
713        let parsed: Value = serde_json::from_str(&result).unwrap();
714        assert_eq!(parsed["status"], "unknown");
715    }
716
717    #[tokio::test]
718    async fn run_background_with_cwd() {
719        let tmp = TempDir::new().unwrap();
720        let sub = tmp.path().join("subdir");
721        std::fs::create_dir(&sub).unwrap();
722        std::fs::write(sub.join("marker.txt"), "content").unwrap();
723
724        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
725        let run_tool = RunBackground::new(tmp.path(), manager.clone());
726        let check_tool = CheckBackground::new(manager.clone());
727
728        let result = run_tool
729            .execute(json!({"command": "ls", "cwd": "subdir"}))
730            .await
731            .unwrap();
732
733        let parsed: Value = serde_json::from_str(&result).unwrap();
734        let job_id = parsed["job_id"].as_str().unwrap().to_string();
735
736        let parsed = poll_until_done(&check_tool, &job_id, Duration::from_secs(5)).await;
737        assert_eq!(parsed["status"], "completed");
738        assert!(parsed["stdout"].as_str().unwrap().contains("marker.txt"));
739    }
740
741    #[tokio::test]
742    async fn run_background_with_env() {
743        let tmp = TempDir::new().unwrap();
744        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
745        let run_tool = RunBackground::new(tmp.path(), manager.clone());
746        let check_tool = CheckBackground::new(manager.clone());
747
748        let result = run_tool
749            .execute(json!({"command": "echo $MY_VAR", "env": {"MY_VAR": "hello"}}))
750            .await
751            .unwrap();
752
753        let parsed: Value = serde_json::from_str(&result).unwrap();
754        let job_id = parsed["job_id"].as_str().unwrap().to_string();
755
756        let parsed = poll_until_done(&check_tool, &job_id, Duration::from_secs(5)).await;
757        assert_eq!(parsed["status"], "completed");
758        assert!(parsed["stdout"].as_str().unwrap().contains("hello"));
759    }
760
761    #[tokio::test]
762    async fn run_background_missing_command() {
763        let tmp = TempDir::new().unwrap();
764        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
765        let run_tool = RunBackground::new(tmp.path(), manager);
766
767        let err = run_tool.execute(json!({})).await.unwrap_err();
768        assert!(matches!(err, Error::BadToolArgs { .. }));
769    }
770
771    #[tokio::test]
772    async fn check_background_missing_job_id() {
773        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
774        let check_tool = CheckBackground::new(manager);
775
776        let err = check_tool.execute(json!({})).await.unwrap_err();
777        assert!(matches!(err, Error::BadToolArgs { .. }));
778    }
779
780    #[tokio::test]
781    async fn take_completed_returns_finished_job() {
782        let mut manager = BackgroundJobManager::new();
783        // Insert a running job
784        let running_id = manager.insert(Job {
785            state: JobState::Running,
786            created_at: Instant::now(),
787        });
788        // Insert a completed job
789        let completed_id = manager.insert(Job {
790            state: JobState::Completed {
791                stdout: "hello".into(),
792                stderr: "".into(),
793                exit_code: 0,
794            },
795            created_at: Instant::now(),
796        });
797
798        // take_completed should return the completed job
799        let (id, output) = manager.take_completed().unwrap();
800        assert_eq!(id, completed_id);
801        assert!(output.contains("hello"));
802        assert!(output.contains("exit code: 0"));
803
804        // Running job should still be there
805        assert!(manager.get_state(&running_id).is_some());
806        // Completed job should be removed
807        assert!(manager.get_state(&completed_id).is_none());
808
809        // No more completed jobs
810        assert!(manager.take_completed().is_none());
811    }
812
813    #[tokio::test]
814    async fn take_completed_returns_failed_job() {
815        let mut manager = BackgroundJobManager::new();
816        manager.insert(Job {
817            state: JobState::Failed {
818                message: "something went wrong".into(),
819            },
820            created_at: Instant::now(),
821        });
822
823        let (id, output) = manager.take_completed().unwrap();
824        assert!(id.starts_with("bg-"));
825        assert!(output.contains("failed"));
826        assert!(output.contains("something went wrong"));
827    }
828
829    #[tokio::test]
830    async fn take_completed_returns_timed_out_job() {
831        let mut manager = BackgroundJobManager::new();
832        manager.insert(Job {
833            state: JobState::TimedOut,
834            created_at: Instant::now(),
835        });
836
837        let (id, output) = manager.take_completed().unwrap();
838        assert!(id.starts_with("bg-"));
839        assert!(output.contains("timed out"));
840    }
841
842    #[tokio::test]
843    async fn take_completed_empty_when_no_finished_jobs() {
844        let mut manager = BackgroundJobManager::new();
845        manager.insert(Job {
846            state: JobState::Running,
847            created_at: Instant::now(),
848        });
849        assert!(manager.take_completed().is_none());
850    }
851
852    #[tokio::test]
853    async fn completed_notify_wakes_on_terminal_state() {
854        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
855
856        let notify = manager.lock().await.completed_notify();
857
858        // notify_one() stores a permit, so a waiter registered AFTER the
859        // state change still wakes — this models the arbiter polling
860        // between turns after a job completed mid-turn (the real usage).
861        {
862            let mut mgr = manager.lock().await;
863            let id = mgr.insert(Job {
864                state: JobState::Running,
865                created_at: Instant::now(),
866            });
867            mgr.update(
868                &id,
869                JobState::Completed {
870                    stdout: "done".into(),
871                    stderr: "".into(),
872                    exit_code: 0,
873                },
874            );
875        }
876
877        let result = tokio::time::timeout(Duration::from_secs(5), notify.notified()).await;
878        assert!(result.is_ok(), "notify should fire on terminal state");
879    }
880
881    #[tokio::test]
882    async fn completed_notify_does_not_fire_for_running_job() {
883        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
884
885        let notify = manager.lock().await.completed_notify();
886
887        // Insert a running job (should NOT trigger notify).
888        {
889            let mut mgr = manager.lock().await;
890            mgr.insert(Job {
891                state: JobState::Running,
892                created_at: Instant::now(),
893            });
894        }
895
896        // The notify should timeout (running jobs don't fire it).
897        let result = tokio::time::timeout(Duration::from_millis(200), notify.notified()).await;
898        assert!(result.is_err(), "notify should not fire for running jobs");
899    }
900
901    #[tokio::test]
902    async fn completed_notify_fires_for_timed_out_job() {
903        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
904
905        let notify = manager.lock().await.completed_notify();
906
907        let handle = tokio::spawn(async move {
908            tokio::time::timeout(Duration::from_secs(5), notify.notified()).await
909        });
910
911        {
912            let mut mgr = manager.lock().await;
913            let id = mgr.insert(Job {
914                state: JobState::Running,
915                created_at: Instant::now(),
916            });
917            mgr.update(&id, JobState::TimedOut);
918        }
919
920        let result = handle.await.unwrap();
921        assert!(result.is_ok(), "notify should fire for TimedOut");
922    }
923
924    #[tokio::test]
925    async fn completed_notify_fires_for_failed_job() {
926        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
927
928        let notify = manager.lock().await.completed_notify();
929
930        let handle = tokio::spawn(async move {
931            tokio::time::timeout(Duration::from_secs(5), notify.notified()).await
932        });
933
934        {
935            let mut mgr = manager.lock().await;
936            let id = mgr.insert(Job {
937                state: JobState::Running,
938                created_at: Instant::now(),
939            });
940            mgr.update(
941                &id,
942                JobState::Failed {
943                    message: "boom".into(),
944                },
945            );
946        }
947
948        let result = handle.await.unwrap();
949        assert!(result.is_ok(), "notify should fire for Failed");
950    }
951
952    #[tokio::test]
953    async fn two_completed_jobs_each_get_a_wake() {
954        // goal-379: `update()` uses notify_one(), which stores a SINGLE
955        // permit. When two jobs reach a terminal state before any waiter
956        // polls, the two notify_one() calls coalesce into one wake. The
957        // drainer consumes that wake, takes ONE job, and — without a re-arm —
958        // the second job would be orphaned forever (the loop idles even
959        // though a background job is finished). take_completed() must re-arm
960        // the notify while terminal jobs remain, so each completed job gets
961        // its own wake.
962        let mut manager = BackgroundJobManager::new();
963        let id_a = manager.insert(Job {
964            state: JobState::Running,
965            created_at: Instant::now(),
966        });
967        let id_b = manager.insert(Job {
968            state: JobState::Running,
969            created_at: Instant::now(),
970        });
971        manager.update(
972            &id_a,
973            JobState::Completed {
974                stdout: "a".into(),
975                stderr: "".into(),
976                exit_code: 0,
977            },
978        );
979        manager.update(
980            &id_b,
981            JobState::Completed {
982                stdout: "b".into(),
983                stderr: "".into(),
984                exit_code: 0,
985            },
986        );
987
988        let notify = manager.completed_notify();
989
990        // First wake consumes the single coalesced permit.
991        tokio::time::timeout(Duration::from_secs(5), notify.notified())
992            .await
993            .expect("first completed job must produce a wake");
994
995        // Drain one job; the other is still terminal, so take_completed must
996        // re-arm the notify for it.
997        let (first_id, first_output) = manager.take_completed().expect("first completed job");
998        assert!(first_output.contains("exit code: 0"));
999
1000        // Second wake must fire for the remaining job (would time out on the
1001        // pre-fix code, where the coalesced permit was already consumed).
1002        tokio::time::timeout(Duration::from_secs(5), notify.notified())
1003            .await
1004            .expect("second completed job must also produce a wake");
1005        let (second_id, second_output) = manager.take_completed().expect("second completed job");
1006        assert!(second_output.contains("exit code: 0"));
1007        assert_ne!(first_id, second_id, "each job must be drained exactly once");
1008
1009        // Queue fully drained; nothing left to take.
1010        assert!(manager.take_completed().is_none());
1011    }
1012
1013    #[tokio::test]
1014    async fn shared_bg_manager_sees_jobs_from_tool() {
1015        let tmp = TempDir::new().unwrap();
1016        let manager = Arc::new(Mutex::new(BackgroundJobManager::new()));
1017
1018        let run_tool = RunBackground::new(tmp.path(), manager.clone());
1019
1020        // Run a quick command
1021        let result = run_tool
1022            .execute(json!({"command": "echo shared-test"}))
1023            .await
1024            .unwrap();
1025
1026        let parsed: Value = serde_json::from_str(&result).unwrap();
1027        let job_id = parsed["job_id"].as_str().unwrap().to_string();
1028        assert_eq!(parsed["status"], "spawned");
1029
1030        // Wait for completion via the shared manager.
1031        let notify = manager.lock().await.completed_notify();
1032        tokio::time::timeout(Duration::from_secs(5), notify.notified())
1033            .await
1034            .expect("notify should fire");
1035
1036        // take_completed from the shared manager should return the job.
1037        let mut mgr = manager.lock().await;
1038        let (id, output) = mgr.take_completed().unwrap();
1039        assert_eq!(id, job_id);
1040        assert!(output.contains("shared-test"));
1041    }
1042
1043    // ── BackgroundJobManager unit tests ─────────────────────────────────────
1044
1045    #[tokio::test]
1046    async fn manager_next_id_increments() {
1047        let mut mgr = BackgroundJobManager::new();
1048        let id0 = mgr.next_id();
1049        let id1 = mgr.next_id();
1050        let id2 = mgr.next_id();
1051        // next_id starts at 1
1052        assert_eq!(id0, "bg-1");
1053        assert_eq!(id1, "bg-2");
1054        assert_eq!(id2, "bg-3");
1055    }
1056
1057    #[tokio::test]
1058    async fn manager_insert_returns_id_and_stores_job() {
1059        let mut mgr = BackgroundJobManager::new();
1060        let job = Job {
1061            state: JobState::Running,
1062            created_at: Instant::now(),
1063        };
1064        let id = mgr.insert(job);
1065        assert_eq!(id, "bg-1");
1066        assert!(mgr.get_state("bg-1").is_some());
1067        assert!(matches!(mgr.get_state("bg-1"), Some(JobState::Running)));
1068    }
1069
1070    #[tokio::test]
1071    async fn manager_get_state_returns_none_for_unknown() {
1072        let mgr = BackgroundJobManager::new();
1073        assert!(mgr.get_state("bg-99").is_none());
1074    }
1075
1076    #[tokio::test]
1077    async fn manager_update_changes_state() {
1078        let mut mgr = BackgroundJobManager::new();
1079        let job = Job {
1080            state: JobState::Running,
1081            created_at: Instant::now(),
1082        };
1083        let id = mgr.insert(job);
1084        mgr.update(
1085            &id,
1086            JobState::Completed {
1087                stdout: "out".into(),
1088                stderr: "".into(),
1089                exit_code: 0,
1090            },
1091        );
1092        assert!(matches!(
1093            mgr.get_state(&id),
1094            Some(JobState::Completed { exit_code: 0, .. })
1095        ));
1096    }
1097
1098    #[tokio::test]
1099    async fn manager_clear_removes_all_jobs() {
1100        let mut mgr = BackgroundJobManager::new();
1101        let id0 = mgr.insert(Job {
1102            state: JobState::Running,
1103            created_at: Instant::now(),
1104        });
1105        let id1 = mgr.insert(Job {
1106            state: JobState::Running,
1107            created_at: Instant::now(),
1108        });
1109        assert!(mgr.get_state(&id0).is_some());
1110        mgr.clear();
1111        assert!(mgr.get_state(&id0).is_none());
1112        assert!(mgr.get_state(&id1).is_none());
1113    }
1114
1115    #[tokio::test]
1116    async fn manager_take_completed_running_job_returns_none() {
1117        let mut mgr = BackgroundJobManager::new();
1118        mgr.insert(Job {
1119            state: JobState::Running,
1120            created_at: Instant::now(),
1121        });
1122        assert!(
1123            mgr.take_completed().is_none(),
1124            "running job must not be taken"
1125        );
1126    }
1127
1128    #[tokio::test]
1129    async fn manager_take_completed_removes_and_returns_job() {
1130        let mut mgr = BackgroundJobManager::new();
1131        let job = Job {
1132            state: JobState::Completed {
1133                stdout: "hello".into(),
1134                stderr: "".into(),
1135                exit_code: 0,
1136            },
1137            created_at: Instant::now(),
1138        };
1139        let id = mgr.insert(job);
1140
1141        let (taken_id, output) = mgr.take_completed().expect("must take completed job");
1142        assert_eq!(taken_id, id);
1143        assert!(output.contains("hello"), "output must contain stdout");
1144        assert!(
1145            output.contains("exit code: 0"),
1146            "output must contain exit code"
1147        );
1148
1149        // After take, job is gone
1150        assert!(mgr.get_state(&id).is_none());
1151    }
1152
1153    #[tokio::test]
1154    async fn manager_take_completed_failed_job() {
1155        let mut mgr = BackgroundJobManager::new();
1156        let job = Job {
1157            state: JobState::Failed {
1158                message: "oops".into(),
1159            },
1160            created_at: Instant::now(),
1161        };
1162        let id = mgr.insert(job);
1163        let (_, output) = mgr.take_completed().expect("must take failed job");
1164        assert!(
1165            output.contains("oops"),
1166            "output must contain failure message"
1167        );
1168        assert!(mgr.get_state(&id).is_none());
1169    }
1170
1171    #[tokio::test]
1172    async fn manager_cleanup_removes_old_jobs() {
1173        let mut mgr = BackgroundJobManager::new();
1174        // Simulate a very old job by using an instant far in the past.
1175        let old_created = Instant::now() - JOB_TTL - Duration::from_secs(1);
1176        let old_job = Job {
1177            state: JobState::Completed {
1178                stdout: "".into(),
1179                stderr: "".into(),
1180                exit_code: 0,
1181            },
1182            created_at: old_created,
1183        };
1184        let old_id = mgr.insert(old_job);
1185
1186        // A fresh running job should survive cleanup.
1187        let fresh_job = Job {
1188            state: JobState::Running,
1189            created_at: Instant::now(),
1190        };
1191        let fresh_id = mgr.insert(fresh_job);
1192
1193        mgr.cleanup();
1194
1195        assert!(
1196            mgr.get_state(&old_id).is_none(),
1197            "old job must be removed by cleanup"
1198        );
1199        assert!(
1200            mgr.get_state(&fresh_id).is_some(),
1201            "fresh job must survive cleanup"
1202        );
1203    }
1204
1205    #[tokio::test]
1206    async fn cleanup_does_not_panic_on_young_process() {
1207        let mut mgr = BackgroundJobManager::new();
1208        // Insert a freshly-created job (created_at = now). A job created in
1209        // this process can never be older than JOB_TTL, so cleanup must keep
1210        // it — and it must not panic even when `now - JOB_TTL` would underflow
1211        // (the case where the Instant epoch is younger than the TTL, e.g.
1212        // shortly after system boot).
1213        let id = mgr.insert(Job {
1214            state: JobState::Running,
1215            created_at: Instant::now(),
1216        });
1217
1218        // Must NOT panic even though `now - JOB_TTL` underflows on a young
1219        // epoch.
1220        mgr.cleanup();
1221
1222        assert!(
1223            mgr.get_state(&id).is_some(),
1224            "freshly-created job must survive cleanup"
1225        );
1226    }
1227}