1use std::collections::HashMap;
6use std::io;
7use std::path::PathBuf;
8use std::sync::atomic::{AtomicU64, Ordering};
9use std::future::Future;
10use std::sync::Arc;
11
12use tokio::sync::{oneshot, Mutex};
13use tokio::task::JoinHandle;
14
15use super::stream::BoundedStream;
16use crate::interpreter::ExecResult;
17
18pub use kaish_types::{JobId, JobInfo, JobStatus};
20
21pub struct Job {
23 pub id: JobId,
25 session_id: u64,
29 pub command: String,
31 handle: Option<JoinHandle<ExecResult>>,
33 result_rx: Option<oneshot::Receiver<ExecResult>>,
35 result: Option<ExecResult>,
37 output_file: Option<PathBuf>,
39 persist_output: bool,
47 stdout_stream: Option<Arc<BoundedStream>>,
49 stderr_stream: Option<Arc<BoundedStream>>,
51 pid: Option<u32>,
53 pgid: Option<u32>,
55 stopped: bool,
57 cancel: Option<tokio_util::sync::CancellationToken>,
63 pgids: Vec<u32>,
68}
69
70impl Job {
71 pub fn new(id: JobId, session_id: u64, command: String, handle: JoinHandle<ExecResult>) -> Self {
73 Self {
74 id,
75 session_id,
76 command,
77 handle: Some(handle),
78 result_rx: None,
79 result: None,
80 output_file: None,
81 persist_output: true,
82 stdout_stream: None,
83 stderr_stream: None,
84 pid: None,
85 pgid: None,
86 stopped: false,
87 cancel: None,
88 pgids: Vec::new(),
89 }
90 }
91
92 pub fn from_channel(id: JobId, session_id: u64, command: String, rx: oneshot::Receiver<ExecResult>) -> Self {
94 Self {
95 id,
96 session_id,
97 command,
98 handle: None,
99 result_rx: Some(rx),
100 result: None,
101 output_file: None,
102 persist_output: true,
103 stdout_stream: None,
104 stderr_stream: None,
105 pid: None,
106 pgid: None,
107 stopped: false,
108 cancel: None,
109 pgids: Vec::new(),
110 }
111 }
112
113 pub fn with_streams(
117 id: JobId,
118 session_id: u64,
119 command: String,
120 rx: oneshot::Receiver<ExecResult>,
121 stdout: Arc<BoundedStream>,
122 stderr: Arc<BoundedStream>,
123 ) -> Self {
124 Self {
125 id,
126 session_id,
127 command,
128 handle: None,
129 result_rx: Some(rx),
130 result: None,
131 output_file: None,
132 persist_output: true,
133 stdout_stream: Some(stdout),
134 stderr_stream: Some(stderr),
135 pid: None,
136 pgid: None,
137 stopped: false,
138 cancel: None,
139 pgids: Vec::new(),
140 }
141 }
142
143 pub fn stopped(id: JobId, session_id: u64, command: String, pid: u32, pgid: u32) -> Self {
145 Self {
146 id,
147 session_id,
148 command,
149 handle: None,
150 result_rx: None,
151 result: None,
152 output_file: None,
153 persist_output: true,
154 stdout_stream: None,
155 stderr_stream: None,
156 pid: Some(pid),
157 pgid: Some(pgid),
158 stopped: true,
159 cancel: None,
160 pgids: Vec::new(),
161 }
162 }
163
164 pub fn output_file(&self) -> Option<&PathBuf> {
166 self.output_file.as_ref()
167 }
168
169 pub fn is_done(&mut self) -> bool {
173 if self.stopped {
174 return false;
175 }
176 self.try_poll();
177 self.result.is_some()
178 }
179
180 pub fn status(&mut self) -> JobStatus {
182 if self.stopped {
183 return JobStatus::Stopped;
184 }
185 self.try_poll();
186 match &self.result {
187 Some(r) if r.ok() => JobStatus::Done,
188 Some(r) if r.latch_request().is_some() => JobStatus::Latched,
192 Some(_) => JobStatus::Failed,
193 None => JobStatus::Running,
194 }
195 }
196
197 pub fn status_string(&mut self) -> String {
205 self.try_poll();
206 match &self.result {
207 Some(r) if r.ok() => "done:0".to_string(),
208 Some(r) if r.latch_request().is_some() => "latched".to_string(),
209 Some(r) => format!("failed:{}", r.code),
210 None => "running".to_string(),
211 }
212 }
213
214 pub fn latch(&mut self) -> Option<kaish_types::result::LatchRequest> {
224 self.try_poll();
225 let id = self.id;
226 self.result.as_ref().and_then(|r| r.latch_request()).map(|mut lr| {
227 lr.job_id = Some(id.0);
228 lr
229 })
230 }
231
232 pub fn stdout_stream(&self) -> Option<&Arc<BoundedStream>> {
234 self.stdout_stream.as_ref()
235 }
236
237 pub fn stderr_stream(&self) -> Option<&Arc<BoundedStream>> {
239 self.stderr_stream.as_ref()
240 }
241
242 fn write_output_file(&self, result: &ExecResult) -> Option<PathBuf> {
244 let is_bytes = result.is_bytes();
248 let text = if is_bytes {
249 std::borrow::Cow::Borrowed("")
250 } else {
251 result.text_out()
252 };
253 if !is_bytes && text.is_empty() && result.err.is_empty() {
254 return None;
255 }
256
257 let tmp_dir = std::env::temp_dir().join("kaish").join("jobs");
258 if std::fs::create_dir_all(&tmp_dir).is_err() {
259 tracing::warn!("Failed to create job output directory");
260 return None;
261 }
262
263 let filename = format!(
272 "session_{}_job_{}.{}.txt",
273 self.session_id,
274 self.id.0,
275 std::process::id()
276 );
277 let path = tmp_dir.join(filename);
278
279 let mut content = String::new();
280 content.push_str(&format!("# Job {}: {}\n", self.id, self.command));
281 content.push_str(&format!("# Status: {}\n\n", if result.ok() { "Done" } else { "Failed" }));
282
283 if is_bytes {
284 let n = result.out_bytes().map(|b| b.len()).unwrap_or(0);
285 content.push_str(&format!(
286 "## STDOUT\n[binary output: {n} bytes — omitted from this text log]\n"
287 ));
288 } else if !text.is_empty() {
289 content.push_str("## STDOUT\n");
290 content.push_str(&text);
291 if !text.ends_with('\n') {
292 content.push('\n');
293 }
294 }
295
296 if !result.err.is_empty() {
297 content.push_str("\n## STDERR\n");
298 content.push_str(&result.err);
299 if !result.err.ends_with('\n') {
300 content.push('\n');
301 }
302 }
303
304 match std::fs::write(&path, content) {
305 Ok(()) => Some(path),
306 Err(e) => {
307 tracing::warn!("Failed to write job output file: {}", e);
308 None
309 }
310 }
311 }
312
313 pub fn cleanup_files(&mut self) {
315 if let Some(path) = self.output_file.take() {
316 if let Err(e) = std::fs::remove_file(&path) {
317 if e.kind() != io::ErrorKind::NotFound {
319 tracing::warn!("Failed to clean up job output file {}: {}", path.display(), e);
320 }
321 }
322 }
323 }
324
325 pub fn try_result(&self) -> Option<&ExecResult> {
327 self.result.as_ref()
328 }
329
330 pub fn try_poll(&mut self) -> bool {
335 if self.result.is_some() {
336 return true;
337 }
338
339 if let Some(rx) = self.result_rx.as_mut() {
341 match rx.try_recv() {
342 Ok(result) => {
343 self.result = Some(result);
344 self.result_rx = None;
345 return true;
346 }
347 Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {
348 return false;
350 }
351 Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
352 self.result = Some(ExecResult::failure(1, "job channel closed"));
354 self.result_rx = None;
355 return true;
356 }
357 }
358 }
359
360 if let Some(handle) = self.handle.as_mut()
362 && handle.is_finished() {
363 let Some(mut handle) = self.handle.take() else {
365 return false;
366 };
367 let waker = std::task::Waker::noop();
369 let mut cx = std::task::Context::from_waker(waker);
370 let result = match std::pin::Pin::new(&mut handle).poll(&mut cx) {
371 std::task::Poll::Ready(Ok(r)) => r,
372 std::task::Poll::Ready(Err(e)) => {
373 ExecResult::failure(1, format!("job panicked: {}", e))
374 }
375 std::task::Poll::Pending => {
376 self.handle = Some(handle);
382 return false;
383 }
384 };
385 self.result = Some(result);
386 return true;
387 }
388
389 false
390 }
391}
392
393static NEXT_SESSION_ID: AtomicU64 = AtomicU64::new(0);
399
400fn prune_orphaned_job_files() {
412 #[cfg(target_os = "linux")]
414 {
415 let jobs_dir = std::env::temp_dir().join("kaish").join("jobs");
416 let Ok(entries) = std::fs::read_dir(&jobs_dir) else {
417 return; };
419 let current_pid = std::process::id();
420 for entry in entries.flatten() {
421 let name = entry.file_name();
422 let name_str = name.to_string_lossy();
423 let file_pid: Option<u32> = name_str
426 .strip_suffix(".txt")
427 .and_then(|s| s.rsplit_once('.'))
428 .and_then(|(_, pid_str)| pid_str.parse().ok());
429 let Some(pid) = file_pid else {
430 continue; };
432 if pid == current_pid {
433 continue; }
435 if std::path::Path::new(&format!("/proc/{}", pid)).exists() {
437 continue; }
439 let _ = std::fs::remove_file(entry.path());
441 }
442 }
443}
444
445pub struct JobManager {
447 session_id: u64,
449 next_id: AtomicU64,
451 jobs: Arc<Mutex<HashMap<JobId, Job>>>,
453 persist_output_files: std::sync::atomic::AtomicBool,
459}
460
461impl JobManager {
462 pub fn new() -> Self {
477 static PRUNE_ONCE: std::sync::Once = std::sync::Once::new();
482 PRUNE_ONCE.call_once(prune_orphaned_job_files);
483 Self {
484 session_id: NEXT_SESSION_ID.fetch_add(1, Ordering::SeqCst),
485 next_id: AtomicU64::new(1),
486 jobs: Arc::new(Mutex::new(HashMap::new())),
487 persist_output_files: std::sync::atomic::AtomicBool::new(true),
488 }
489 }
490
491 pub fn set_persist_output_files(&self, on: bool) {
501 self.persist_output_files.store(on, Ordering::Relaxed);
502 }
503
504 pub fn persist_output_files(&self) -> bool {
506 self.persist_output_files.load(Ordering::Relaxed)
507 }
508
509 pub async fn spawn<F>(&self, command: String, future: F) -> JobId
514 where
515 F: std::future::Future<Output = ExecResult> + Send + 'static,
516 {
517 let id = JobId(self.next_id.fetch_add(1, Ordering::SeqCst));
518 let handle = tokio::spawn(crate::telemetry::bind_current_context(future));
521 let mut job = Job::new(id, self.session_id, command, handle);
522 job.persist_output = self.persist_output_files();
523
524 self.jobs.lock().await.insert(id, job);
531
532 id
533 }
534
535 pub async fn register(&self, command: String, rx: oneshot::Receiver<ExecResult>) -> JobId {
537 let id = JobId(self.next_id.fetch_add(1, Ordering::SeqCst));
538 let mut job = Job::from_channel(id, self.session_id, command, rx);
539 job.persist_output = self.persist_output_files();
540
541 let mut jobs = self.jobs.lock().await;
542 jobs.insert(id, job);
543
544 id
545 }
546
547 pub async fn register_with_streams(
551 &self,
552 command: String,
553 rx: oneshot::Receiver<ExecResult>,
554 stdout: Arc<BoundedStream>,
555 stderr: Arc<BoundedStream>,
556 ) -> JobId {
557 let id = JobId(self.next_id.fetch_add(1, Ordering::SeqCst));
558 let mut job = Job::with_streams(id, self.session_id, command, rx, stdout, stderr);
559 job.persist_output = self.persist_output_files();
560
561 let mut jobs = self.jobs.lock().await;
562 jobs.insert(id, job);
563
564 id
565 }
566
567 pub async fn wait(&self, id: JobId) -> Option<ExecResult> {
577 loop {
589 {
590 let mut jobs = self.jobs.lock().await;
591 let job = jobs.get_mut(&id)?;
592 if job.is_done() {
593 let result = job
594 .result
595 .clone()
596 .unwrap_or_else(|| ExecResult::failure(1, "no result"));
597 if job.persist_output
599 && job.output_file.is_none()
600 && let Some(path) = job.write_output_file(&result)
601 {
602 job.output_file = Some(path);
603 }
604 return Some(result);
605 }
606 }
607 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
609 }
610 }
611
612 pub async fn wait_all(&self) -> Vec<(JobId, ExecResult)> {
614 let mut results = Vec::new();
615
616 let ids: Vec<JobId> = {
618 let jobs = self.jobs.lock().await;
619 jobs.keys().copied().collect()
620 };
621
622 for id in ids {
623 if let Some(result) = self.wait(id).await {
624 results.push((id, result));
625 }
626 }
627
628 results
629 }
630
631 pub async fn list(&self) -> Vec<JobInfo> {
633 let mut jobs = self.jobs.lock().await;
634 jobs.values_mut()
635 .map(|job| {
636 let status = job.status();
637 let latch = job.latch();
638 JobInfo::new(job.id, job.command.clone(), status)
639 .with_output_file(job.output_file.clone())
640 .with_pid(job.pid)
641 .with_latch(latch)
642 })
643 .collect()
644 }
645
646 pub async fn running_count(&self) -> usize {
648 let mut jobs = self.jobs.lock().await;
649 let mut count = 0;
650 for job in jobs.values_mut() {
651 if !job.is_done() {
652 count += 1;
653 }
654 }
655 count
656 }
657
658 pub async fn reap_finished(&self) -> Vec<JobInfo> {
671 let mut jobs = self.jobs.lock().await;
672 let done_ids: Vec<JobId> = jobs
673 .iter_mut()
674 .filter_map(|(id, job)| (job.is_done() && job.latch().is_none()).then_some(*id))
675 .collect();
676
677 let mut removed = Vec::with_capacity(done_ids.len());
678 for id in done_ids {
679 let Some(mut job) = jobs.remove(&id) else {
680 continue;
681 };
682 let status = job.status();
683 let info = JobInfo::new(job.id, job.command.clone(), status).with_pid(job.pid);
684 job.cleanup_files();
685 removed.push(info);
686 }
687 removed
688 }
689
690 pub async fn cleanup(&self) {
695 self.reap_finished().await;
696 }
697
698 pub async fn exists(&self, id: JobId) -> bool {
700 let jobs = self.jobs.lock().await;
701 jobs.contains_key(&id)
702 }
703
704 pub async fn is_latched(&self, id: JobId) -> bool {
708 let mut jobs = self.jobs.lock().await;
709 jobs.get_mut(&id).is_some_and(|job| job.latch().is_some())
710 }
711
712 pub async fn get(&self, id: JobId) -> Option<JobInfo> {
714 let mut jobs = self.jobs.lock().await;
715 jobs.get_mut(&id).map(|job| {
716 let status = job.status();
717 let latch = job.latch();
718 JobInfo::new(job.id, job.command.clone(), status)
719 .with_output_file(job.output_file.clone())
720 .with_pid(job.pid)
721 .with_latch(latch)
722 })
723 }
724
725 pub async fn get_command(&self, id: JobId) -> Option<String> {
727 let jobs = self.jobs.lock().await;
728 jobs.get(&id).map(|job| job.command.clone())
729 }
730
731 pub async fn get_status_string(&self, id: JobId) -> Option<String> {
733 let mut jobs = self.jobs.lock().await;
734 jobs.get_mut(&id).map(|job| job.status_string())
735 }
736
737 pub async fn get_latch(&self, id: JobId) -> Option<kaish_types::result::LatchRequest> {
742 let mut jobs = self.jobs.lock().await;
743 jobs.get_mut(&id).and_then(|job| job.latch())
744 }
745
746 pub async fn read_stdout(&self, id: JobId) -> Option<Vec<u8>> {
750 let jobs = self.jobs.lock().await;
751 if let Some(job) = jobs.get(&id)
752 && let Some(stream) = job.stdout_stream() {
753 return Some(stream.read().await);
754 }
755 None
756 }
757
758 pub async fn read_stderr(&self, id: JobId) -> Option<Vec<u8>> {
762 let jobs = self.jobs.lock().await;
763 if let Some(job) = jobs.get(&id)
764 && let Some(stream) = job.stderr_stream() {
765 return Some(stream.read().await);
766 }
767 None
768 }
769
770 pub async fn list_ids(&self) -> Vec<JobId> {
772 let jobs = self.jobs.lock().await;
773 jobs.keys().copied().collect()
774 }
775
776 pub async fn register_stopped(&self, command: String, pid: u32, pgid: u32) -> JobId {
778 let id = JobId(self.next_id.fetch_add(1, Ordering::SeqCst));
779 let job = Job::stopped(id, self.session_id, command, pid, pgid);
780 let mut jobs = self.jobs.lock().await;
781 jobs.insert(id, job);
782 id
783 }
784
785 pub async fn stop_job(&self, id: JobId, pid: u32, pgid: u32) {
787 let mut jobs = self.jobs.lock().await;
788 if let Some(job) = jobs.get_mut(&id) {
789 job.stopped = true;
790 job.pid = Some(pid);
791 job.pgid = Some(pgid);
792 }
793 }
794
795 pub async fn resume_job(&self, id: JobId) {
797 let mut jobs = self.jobs.lock().await;
798 if let Some(job) = jobs.get_mut(&id) {
799 job.stopped = false;
800 }
801 }
802
803 pub async fn last_stopped(&self) -> Option<JobId> {
805 let mut jobs = self.jobs.lock().await;
806 let mut best: Option<JobId> = None;
808 for job in jobs.values_mut() {
809 if job.stopped {
810 match best {
811 None => best = Some(job.id),
812 Some(b) if job.id.0 > b.0 => best = Some(job.id),
813 _ => {}
814 }
815 }
816 }
817 best
818 }
819
820 pub async fn get_process_info(&self, id: JobId) -> Option<(u32, u32)> {
822 let jobs = self.jobs.lock().await;
823 jobs.get(&id).and_then(|job| {
824 match (job.pid, job.pgid) {
825 (Some(pid), Some(pgid)) => Some((pid, pgid)),
826 _ => None,
827 }
828 })
829 }
830
831 pub async fn set_cancel_token(&self, id: JobId, token: tokio_util::sync::CancellationToken) {
835 let mut jobs = self.jobs.lock().await;
836 if let Some(job) = jobs.get_mut(&id) {
837 job.cancel = Some(token);
838 }
839 }
840
841 pub async fn cancel(&self, id: JobId) -> bool {
845 let jobs = self.jobs.lock().await;
846 match jobs.get(&id).and_then(|job| job.cancel.clone()) {
847 Some(token) => {
848 token.cancel();
849 true
850 }
851 None => false,
852 }
853 }
854
855 pub async fn add_pgid(&self, id: JobId, pgid: u32) {
859 let mut jobs = self.jobs.lock().await;
860 if let Some(job) = jobs.get_mut(&id) {
861 if !job.pgids.contains(&pgid) {
862 job.pgids.push(pgid);
863 }
864 }
865 }
866
867 pub async fn job_pgids(&self, id: JobId) -> Vec<u32> {
871 let jobs = self.jobs.lock().await;
872 jobs.get(&id)
873 .map(|job| {
874 let mut v = job.pgids.clone();
875 if let Some(pg) = job.pgid {
876 if !v.contains(&pg) {
877 v.push(pg);
878 }
879 }
880 v
881 })
882 .unwrap_or_default()
883 }
884
885 pub async fn remove(&self, id: JobId) {
892 let mut jobs = self.jobs.lock().await;
893 if let Some(mut job) = jobs.remove(&id) {
894 job.cleanup_files();
895 }
896 }
897}
898
899impl Default for JobManager {
900 fn default() -> Self {
901 Self::new()
902 }
903}
904
905#[cfg(test)]
906mod tests {
907 use super::*;
908 use std::time::Duration;
909
910 #[tokio::test]
911 async fn test_no_host_output_file_when_persistence_disabled() {
912 let manager = JobManager::new();
916 assert!(manager.persist_output_files(), "default is to persist");
917 manager.set_persist_output_files(false);
918 assert!(!manager.persist_output_files());
919
920 let id = manager.spawn("leaky".to_string(), async {
921 ExecResult::success("output that must not hit host disk")
922 }).await;
923 tokio::time::sleep(Duration::from_millis(10)).await;
924 let result = manager.wait(id).await;
925 assert!(result.is_some());
926
927 let output_file = {
929 let jobs = manager.jobs.lock().await;
930 jobs.get(&id).and_then(|j| j.output_file().cloned())
931 };
932 assert!(
933 output_file.is_none(),
934 "no host output file should be written when persistence is disabled, got {output_file:?}"
935 );
936 }
937
938 #[tokio::test]
939 async fn test_spawn_and_wait() {
940 let manager = JobManager::new();
941
942 let id = manager.spawn("test".to_string(), async {
943 tokio::time::sleep(Duration::from_millis(10)).await;
944 ExecResult::success("done")
945 }).await;
946
947 tokio::time::sleep(Duration::from_millis(5)).await;
949
950 let result = manager.wait(id).await;
951 assert!(result.is_some());
952 let result = result.unwrap();
953 assert!(result.ok());
954 assert_eq!(&*result.text_out(), "done");
955 }
956
957 #[tokio::test]
958 async fn test_wait_all() {
959 let manager = JobManager::new();
960
961 manager.spawn("job1".to_string(), async {
962 tokio::time::sleep(Duration::from_millis(10)).await;
963 ExecResult::success("one")
964 }).await;
965
966 manager.spawn("job2".to_string(), async {
967 tokio::time::sleep(Duration::from_millis(5)).await;
968 ExecResult::success("two")
969 }).await;
970
971 tokio::time::sleep(Duration::from_millis(5)).await;
973
974 let results = manager.wait_all().await;
975 assert_eq!(results.len(), 2);
976 }
977
978 #[tokio::test]
979 async fn test_list_jobs() {
980 let manager = JobManager::new();
981
982 manager.spawn("test job".to_string(), async {
983 tokio::time::sleep(Duration::from_millis(50)).await;
984 ExecResult::success("")
985 }).await;
986
987 tokio::time::sleep(Duration::from_millis(5)).await;
989
990 let jobs = manager.list().await;
991 assert_eq!(jobs.len(), 1);
992 assert_eq!(jobs[0].command, "test job");
993 assert_eq!(jobs[0].status, JobStatus::Running);
994 }
995
996 #[tokio::test]
997 async fn latch_stamps_job_id_back_reference() {
998 let manager = JobManager::new();
1004
1005 let id = manager.spawn("gated".to_string(), async {
1006 let mut result = ExecResult::failure(2, "confirmation required");
1007 result.latch = Some(Box::new(kaish_types::result::LatchRequest {
1008 nonce: "a3f7b2c1".to_string(),
1009 command: "rm".to_string(),
1010 paths: vec!["x".to_string()],
1011 hint: "rm --confirm=a3f7b2c1 x".to_string(),
1012 tool: "rm".to_string(),
1013 argv: vec!["x".to_string()],
1014 ttl: 60,
1015 job_id: None, }));
1017 result
1018 }).await;
1019
1020 tokio::time::sleep(Duration::from_millis(10)).await;
1021
1022 let latch = manager.get_latch(id).await.expect("job must be latched");
1023 assert_eq!(
1024 latch.job_id,
1025 Some(id.0),
1026 "Job::latch() must stamp this job's own id onto the surfaced request"
1027 );
1028 }
1029
1030 #[tokio::test]
1031 async fn test_job_status_after_completion() {
1032 let manager = JobManager::new();
1033
1034 let id = manager.spawn("quick".to_string(), async {
1035 ExecResult::success("")
1036 }).await;
1037
1038 tokio::time::sleep(Duration::from_millis(10)).await;
1040 let _ = manager.wait(id).await;
1041
1042 let info = manager.get(id).await;
1043 assert!(info.is_some());
1044 assert_eq!(info.unwrap().status, JobStatus::Done);
1045 }
1046
1047 #[tokio::test]
1048 async fn test_cleanup() {
1049 let manager = JobManager::new();
1050
1051 let id = manager.spawn("done".to_string(), async {
1052 ExecResult::success("")
1053 }).await;
1054
1055 tokio::time::sleep(Duration::from_millis(10)).await;
1057 let _ = manager.wait(id).await;
1058
1059 assert_eq!(manager.list().await.len(), 1);
1061
1062 manager.cleanup().await;
1064
1065 assert_eq!(manager.list().await.len(), 0);
1067 }
1068
1069 #[tokio::test]
1070 async fn test_cleanup_removes_temp_files() {
1071 let manager = JobManager::new();
1073
1074 let id = manager.spawn("output job".to_string(), async {
1075 ExecResult::success("some output that gets written to a temp file")
1076 }).await;
1077
1078 tokio::time::sleep(Duration::from_millis(10)).await;
1080 let result = manager.wait(id).await;
1081 assert!(result.is_some());
1082
1083 let output_file = {
1087 let jobs = manager.jobs.lock().await;
1088 jobs.get(&id).and_then(|j| j.output_file().cloned())
1089 };
1090 let path = output_file.expect("job with output should have written a temp file");
1091 assert!(path.exists(), "temp file should exist before cleanup: {}", path.display());
1092
1093 manager.cleanup().await;
1095
1096 assert!(
1097 !path.exists(),
1098 "temp file should be removed after cleanup: {}",
1099 path.display()
1100 );
1101 }
1102
1103 #[tokio::test]
1104 async fn test_reap_finished_returns_removed_job_info() {
1105 let manager = JobManager::new();
1108 manager.set_persist_output_files(false);
1109
1110 let id = manager
1111 .spawn("sleep 0.1".to_string(), async { ExecResult::success("") })
1112 .await;
1113 tokio::time::sleep(Duration::from_millis(10)).await;
1114 let _ = manager.wait(id).await;
1115
1116 let removed = manager.reap_finished().await;
1117 assert_eq!(removed.len(), 1);
1118 assert_eq!(removed[0].id, id);
1119 assert_eq!(removed[0].command, "sleep 0.1");
1120 assert_eq!(removed[0].status, JobStatus::Done);
1121
1122 assert!(manager.list().await.is_empty());
1124 }
1125
1126 #[tokio::test]
1127 async fn test_reap_finished_never_reaps_latched_jobs() {
1128 use kaish_types::result::LatchRequest;
1134
1135 let manager = JobManager::new();
1136 manager.set_persist_output_files(false);
1137 let (tx, rx) = oneshot::channel();
1138 let id = manager.register("rm precious.txt".to_string(), rx).await;
1139
1140 let mut gated = ExecResult::failure(2, "rm: confirmation required (latch enabled)");
1141 gated.latch = Some(Box::new(LatchRequest {
1142 nonce: "a3f7b2c1".to_string(),
1143 command: "rm".to_string(),
1144 paths: vec!["precious.txt".to_string()],
1145 hint: "rm --confirm=\"a3f7b2c1\" precious.txt".to_string(),
1146 tool: "rm".to_string(),
1147 argv: vec!["precious.txt".to_string()],
1148 ttl: 60,
1149 job_id: None,
1150 }));
1151 tx.send(gated).expect("send gated result");
1152 tokio::time::sleep(Duration::from_millis(10)).await;
1153
1154 let info = manager.get(id).await.expect("job exists");
1156 assert_eq!(info.status, JobStatus::Latched);
1157
1158 let removed = manager.reap_finished().await;
1159 assert!(
1160 removed.is_empty(),
1161 "a latched job must never be auto-reaped: {removed:?}"
1162 );
1163 assert_eq!(
1164 manager.list().await.len(),
1165 1,
1166 "the latched job must still be tracked so its gate can be fulfilled"
1167 );
1168 }
1169
1170 #[tokio::test]
1171 async fn test_register_with_channel() {
1172 let manager = JobManager::new();
1173 let (tx, rx) = oneshot::channel();
1174
1175 let id = manager.register("channel job".to_string(), rx).await;
1176
1177 tx.send(ExecResult::success("from channel")).unwrap();
1179
1180 let result = manager.wait(id).await;
1181 assert!(result.is_some());
1182 assert_eq!(&*result.unwrap().text_out(), "from channel");
1183 }
1184
1185 #[tokio::test]
1186 async fn test_spawn_immediately_available() {
1187 let manager = JobManager::new();
1189
1190 let id = manager.spawn("instant".to_string(), async {
1191 tokio::time::sleep(Duration::from_millis(100)).await;
1192 ExecResult::success("done")
1193 }).await;
1194
1195 let exists = manager.exists(id).await;
1197 assert!(exists, "job should be immediately available after spawn()");
1198
1199 let info = manager.get(id).await;
1200 assert!(info.is_some(), "job info should be available immediately");
1201 }
1202
1203 #[tokio::test]
1204 async fn test_nonexistent_job() {
1205 let manager = JobManager::new();
1206 let result = manager.wait(JobId(999)).await;
1207 assert!(result.is_none());
1208 }
1209
1210 #[tokio::test]
1211 async fn test_cancel_token_fires() {
1212 let manager = JobManager::new();
1215 let token = tokio_util::sync::CancellationToken::new();
1216 let id = manager.spawn("bg".to_string(), async { ExecResult::success("") }).await;
1217 manager.set_cancel_token(id, token.clone()).await;
1218
1219 assert!(!token.is_cancelled());
1220 assert!(manager.cancel(id).await, "cancel should report success");
1221 assert!(token.is_cancelled(), "the job's token must be tripped");
1222 }
1223
1224 #[tokio::test]
1225 async fn test_cancel_without_token_returns_false() {
1226 let manager = JobManager::new();
1227 let id = manager.spawn("bg".to_string(), async { ExecResult::success("") }).await;
1228 assert!(!manager.cancel(id).await);
1230 assert!(!manager.cancel(JobId(999)).await);
1232 }
1233
1234 #[tokio::test]
1235 async fn test_pgids_recorded_and_deduped() {
1236 let manager = JobManager::new();
1237 let id = manager.spawn("bg".to_string(), async { ExecResult::success("") }).await;
1238 assert!(manager.job_pgids(id).await.is_empty());
1239
1240 manager.add_pgid(id, 4242).await;
1241 manager.add_pgid(id, 4243).await;
1242 manager.add_pgid(id, 4242).await; assert_eq!(manager.job_pgids(id).await, vec![4242, 4243]);
1244
1245 assert!(manager.job_pgids(JobId(999)).await.is_empty());
1247 }
1248
1249 #[tokio::test]
1250 async fn wait_does_not_block_other_job_ops() {
1251 let manager = Arc::new(JobManager::new());
1258 manager.set_persist_output_files(false);
1259
1260 let (tx, rx) = oneshot::channel::<()>();
1262 let id = manager
1263 .spawn("blocker".to_string(), async move {
1264 let _ = rx.await;
1265 ExecResult::success("done")
1266 })
1267 .await;
1268
1269 let waiter = {
1272 let m = manager.clone();
1273 tokio::spawn(async move { m.wait(id).await })
1274 };
1275 tokio::time::sleep(Duration::from_millis(50)).await;
1277
1278 let listed = tokio::time::timeout(Duration::from_secs(2), manager.list()).await;
1280 assert!(
1281 listed.is_ok(),
1282 "list() blocked while wait() was parked — jobs lock held across await"
1283 );
1284 let second = tokio::time::timeout(
1285 Duration::from_secs(2),
1286 manager.spawn("second".to_string(), async { ExecResult::success("2") }),
1287 )
1288 .await;
1289 assert!(
1290 second.is_ok(),
1291 "spawn() blocked/spun while wait() was parked"
1292 );
1293
1294 let _ = tx.send(());
1296 let result = tokio::time::timeout(Duration::from_secs(2), waiter)
1297 .await
1298 .expect("waiter join timed out")
1299 .expect("waiter task panicked");
1300 assert_eq!(result.map(|r| r.code), Some(0), "waiter should see exit 0");
1301 }
1302
1303 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1304 async fn wait_survives_a_dropped_waiter() {
1305 let manager = Arc::new(JobManager::new());
1312 manager.set_persist_output_files(false);
1313
1314 let (tx, rx) = oneshot::channel::<()>();
1315 let id = manager
1316 .spawn("blocker".to_string(), async move {
1317 let _ = rx.await;
1318 ExecResult::success("done")
1319 })
1320 .await;
1321
1322 {
1324 let m = manager.clone();
1325 let a = tokio::spawn(async move { m.wait(id).await });
1326 tokio::time::sleep(Duration::from_millis(20)).await;
1327 a.abort();
1328 let _ = a.await;
1329 }
1330
1331 let _ = tx.send(());
1333
1334 let res = tokio::time::timeout(Duration::from_secs(2), manager.wait(id))
1336 .await
1337 .expect("wait must not hang after a prior waiter was dropped");
1338 assert_eq!(res.map(|r| r.code), Some(0), "B should see the completed job");
1339 }
1340}