kaish-kernel 0.16.0

//! wait — Wait for background jobs to complete.

use async_trait::async_trait;
use clap::{CommandFactory, Parser};


use crate::ast::Value;
use crate::interpreter::{ExecResult, OutputData};
use crate::scheduler::JobId;
use crate::tools::{schema_from_clap, ExecContext, ToolCtx, GlobalFlags, Tool, ToolArgs, ToolSchema};

/// Wait tool: wait for background jobs.
pub struct Wait;

/// clap-derived argv layer for wait.
#[derive(Parser, Debug)]
#[command(name = "wait", about = "Wait for background jobs to complete")]
struct WaitArgs {
    #[command(flatten)]
    global: GlobalFlags,

    /// Job id, with or without the `%` prefix (`%1` or `1`). Waits for every
    /// background job when omitted. A non-numeric id exits 1. A stopped job
    /// cannot finish: waiting on one exits 1 — resume it with `bg` or `fg`
    /// first.
    job: Vec<String>,
}

#[async_trait]
impl Tool for Wait {
    fn name(&self) -> &str {
        "wait"
    }

    fn schema(&self) -> ToolSchema {
        schema_from_clap(
            &WaitArgs::command(),
            "wait",
            "Wait for background jobs to complete",
            [
                ("Wait for all jobs", "wait"),
                ("Wait for specific job", "wait 1"),
            ],
        )
    }

    async fn execute(&self, args: ToolArgs, ctx: &mut dyn ToolCtx) -> ExecResult {
        let Some(ctx) = ctx.as_any_mut().downcast_mut::<ExecContext>() else {
            return ExecResult::failure(1, "internal error: kernel builtin requires ExecContext");
        };
        let argv = match args.to_argv() {
            Ok(v) => v,
            Err(e) => return ExecResult::failure(2, format!("wait: {e}")),
        };
        let parsed = match WaitArgs::try_parse_from(
            std::iter::once("wait".to_string()).chain(argv),
        ) {
            Ok(p) => p,
            Err(e) => return ExecResult::failure(2, format!("wait: {e}")),
        };
        parsed.global.apply(ctx);

        let manager = match &ctx.job_manager {
            Some(m) => m.clone(),
            None => return ExecResult::with_output(OutputData::text("(no job manager)\n")),
        };

        // Wait for the specific job ids provided (`wait 1`, `wait %1 %2`).
        if !args.positional.is_empty() {
            let mut output = String::new();
            let mut any_failed = false;

            for spec in &args.positional {
                let id = match spec {
                    Value::Int(i) => JobId(*i as u64),
                    Value::String(s) => {
                        // Accept the bash jobspec form `%N` as well as a bare
                        // number; the `%` is a job marker, not part of the id.
                        let digits = s.strip_prefix('%').unwrap_or(s);
                        match digits.parse::<u64>() {
                            Ok(i) => JobId(i),
                            Err(_) => return ExecResult::failure(1, format!("wait: invalid job id: {}", s)),
                        }
                    }
                    _ => return ExecResult::failure(1, "wait: job id must be a number"),
                };

                match manager.wait(id).await {
                    Some(result) => {
                        let killed = matches!(
                            manager.get(id).await.map(|info| info.status),
                            Some(crate::scheduler::JobStatus::Killed)
                        );
                        let status = classify(&result, killed, &mut any_failed);
                        output.push_str(&format!("[{}] {}\n", id, status));
                    }
                    // `wait` returns None for a missing job AND for a stopped
                    // one (a stopped job can never finish; polling it hung
                    // forever). Tell them apart for the error text.
                    None => {
                        if manager.get(id).await.is_some() {
                            return ExecResult::failure(
                                1,
                                format!(
                                    "wait: job {} is stopped and cannot finish — resume it with bg or fg, then wait",
                                    id
                                ),
                            );
                        }
                        return ExecResult::failure(1, format!("wait: job {} not found", id));
                    }
                }
            }

            finish(output, any_failed)
        } else {
            // Wait for all jobs
            let results = manager.wait_all().await;

            if results.is_empty() {
                return ExecResult::with_output(OutputData::text("(no jobs to wait for)\n"));
            }

            let mut output = String::new();
            let mut any_failed = false;

            for (id, result) in results {
                let killed = matches!(
                    manager.get(id).await.map(|info| info.status),
                    Some(crate::scheduler::JobStatus::Killed)
                );
                let status = classify(&result, killed, &mut any_failed);
                output.push_str(&format!("[{}] {}\n", id, status));
            }

            finish(output, any_failed)
        }
    }
}

/// Classify one waited job's result into a display word, threading the
/// aggregate `any_failed` flag.
fn classify(result: &ExecResult, killed: bool, any_failed: &mut bool) -> &'static str {
    if result.ok() {
        "Done"
    } else {
        *any_failed = true;
        // A kill is not an organic failure — say so (GH #244). The exit code
        // still aggregates as a failure: the waited-on work did not finish.
        if killed {
            "Killed"
        } else {
            "Failed"
        }
    }
}

/// Assemble `wait`'s result: any failure is exit 1; otherwise success.
fn finish(output: String, any_failed: bool) -> ExecResult {
    if any_failed {
        let mut result = ExecResult::from_output(1, output.clone(), "");
        result.set_output(Some(OutputData::text(output)));
        result
    } else {
        ExecResult::with_output(OutputData::text(output))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::scheduler::JobManager;
    use crate::vfs::{MemoryFs, VfsRouter};
    use std::sync::Arc;
    use std::time::Duration;

    fn make_ctx() -> ExecContext {
        let mut vfs = VfsRouter::new();
        vfs.mount("/", MemoryFs::new());
        ExecContext::new(Arc::new(vfs))
    }

    #[tokio::test]
    async fn test_wait_no_manager() {
        let mut ctx = make_ctx();
        let result = Wait.execute(ToolArgs::new(), &mut ctx).await;
        assert!(result.ok());
        assert!(result.text_out().contains("no job manager"));
    }

    #[tokio::test]
    async fn test_wait_no_jobs() {
        let mut ctx = make_ctx();
        ctx.set_job_manager(Arc::new(JobManager::new()));

        let result = Wait.execute(ToolArgs::new(), &mut ctx).await;
        assert!(result.ok());
        assert!(result.text_out().contains("no jobs"));
    }

    #[tokio::test]
    async fn test_wait_all() {
        let mut ctx = make_ctx();
        let manager = Arc::new(JobManager::new());
        ctx.set_job_manager(manager.clone());

        // Spawn jobs
        manager.spawn("job1".to_string(), async {
            tokio::time::sleep(Duration::from_millis(10)).await;
            ExecResult::success("")
        }).await;
        manager.spawn("job2".to_string(), async {
            tokio::time::sleep(Duration::from_millis(5)).await;
            ExecResult::success("")
        }).await;

        // Wait for jobs to register
        tokio::time::sleep(Duration::from_millis(5)).await;

        let result = Wait.execute(ToolArgs::new(), &mut ctx).await;
        assert!(result.ok());
        assert!(result.text_out().contains("[1]"));
        assert!(result.text_out().contains("[2]"));
        assert!(result.text_out().contains("Done"));
    }

    #[tokio::test]
    async fn test_wait_specific_job() {
        let mut ctx = make_ctx();
        let manager = Arc::new(JobManager::new());
        ctx.set_job_manager(manager.clone());

        // Spawn a job
        let id = manager.spawn("test".to_string(), async {
            tokio::time::sleep(Duration::from_millis(10)).await;
            ExecResult::success("")
        }).await;

        // Wait for job to register
        tokio::time::sleep(Duration::from_millis(5)).await;

        let mut args = ToolArgs::new();
        args.positional.push(Value::Int(id.0 as i64));

        let result = Wait.execute(args, &mut ctx).await;
        assert!(result.ok());
        assert!(result.text_out().contains(&format!("[{}]", id)));
    }

    #[tokio::test]
    async fn test_wait_jobspec_percent_form() {
        // `wait %1` — the `%` is a job marker, stripped before lookup.
        let mut ctx = make_ctx();
        let manager = Arc::new(JobManager::new());
        ctx.set_job_manager(manager.clone());

        let id = manager.spawn("test".to_string(), async {
            tokio::time::sleep(Duration::from_millis(10)).await;
            ExecResult::success("")
        }).await;
        tokio::time::sleep(Duration::from_millis(5)).await;

        let mut args = ToolArgs::new();
        args.positional.push(Value::String(format!("%{}", id.0)));

        let result = Wait.execute(args, &mut ctx).await;
        assert!(result.ok());
        assert!(result.text_out().contains(&format!("[{}]", id)));
    }

    #[tokio::test]
    async fn test_wait_multiple_jobspecs() {
        // `wait %1 %2` waits for each named job, not just the first.
        let mut ctx = make_ctx();
        let manager = Arc::new(JobManager::new());
        ctx.set_job_manager(manager.clone());

        let id1 = manager.spawn("j1".to_string(), async {
            tokio::time::sleep(Duration::from_millis(10)).await;
            ExecResult::success("")
        }).await;
        let id2 = manager.spawn("j2".to_string(), async {
            tokio::time::sleep(Duration::from_millis(5)).await;
            ExecResult::success("")
        }).await;
        tokio::time::sleep(Duration::from_millis(5)).await;

        let mut args = ToolArgs::new();
        args.positional.push(Value::String(format!("%{}", id1.0)));
        args.positional.push(Value::String(format!("%{}", id2.0)));

        let result = Wait.execute(args, &mut ctx).await;
        assert!(result.ok());
        assert!(result.text_out().contains(&format!("[{}]", id1)));
        assert!(result.text_out().contains(&format!("[{}]", id2)));
    }

    #[tokio::test]
    async fn test_wait_nonexistent_job() {
        let mut ctx = make_ctx();
        ctx.set_job_manager(Arc::new(JobManager::new()));

        let mut args = ToolArgs::new();
        args.positional.push(Value::Int(999));

        let result = Wait.execute(args, &mut ctx).await;
        assert!(!result.ok());
        assert!(result.err.contains("not found"));
    }

    #[tokio::test]
    async fn test_wait_failed_job() {
        let mut ctx = make_ctx();
        let manager = Arc::new(JobManager::new());
        ctx.set_job_manager(manager.clone());

        // Spawn a failing job
        manager.spawn("fail".to_string(), async {
            ExecResult::failure(1, "intentional failure")
        }).await;

        // Wait for job to register
        tokio::time::sleep(Duration::from_millis(5)).await;

        let result = Wait.execute(ToolArgs::new(), &mut ctx).await;
        assert!(!result.ok()); // Overall result fails if any job failed
        assert!(result.text_out().contains("Failed"));
    }

    /// `wait %N` on a stopped job errors with the resume instruction, not
    /// "not found" — the builtin distinguishes None-because-stopped from
    /// None-because-missing.
    #[tokio::test]
    async fn test_wait_on_stopped_job_gives_actionable_error() {
        let mut ctx = make_ctx();
        let manager = Arc::new(JobManager::new());
        ctx.set_job_manager(manager.clone());

        let id = manager
            .register_stopped("sleep 30".to_string(), 4242, 4242)
            .await;

        let mut args = ToolArgs::new();
        args.positional.push(Value::Int(id.0 as i64));

        let result = Wait.execute(args, &mut ctx).await;
        assert!(!result.ok(), "waiting on a stopped job must fail loud");
        assert!(
            result.err.contains("stopped and cannot finish"),
            "error must say the job is stopped: {}",
            result.err
        );
        assert!(
            result.err.contains("bg or fg"),
            "error must say how to resume: {}",
            result.err
        );
    }
}