1use std::collections::HashMap;
66use std::io::Read;
67use std::path::PathBuf;
68use std::sync::{Arc, Mutex};
69use std::time::{Duration, SystemTime, UNIX_EPOCH};
70
71use portable_pty::{native_pty_system, CommandBuilder, PtySize};
72use thiserror::Error;
73use tokio::sync::{mpsc, oneshot};
74use tokio::task;
75
76use crate::beholders::{registry_with_user_beholders, BeholderSelect};
77use crate::store::{RunFilter, StoreError, TaskStore};
78use crate::types::{BeholderStatus, Initiator, OutputChunk, RunStatus, Stream, TaskRunId, TaskRunMeta};
79
80const DEFAULT_GRACE: Duration = Duration::from_secs(5);
81const READ_BUF_SIZE: usize = 4096;
82const SIGTERM: i32 = 15;
83const SIGKILL: i32 = 9;
84
85#[derive(Debug, Error)]
88pub enum DriverError {
89 #[error("store: {0}")]
90 Store(#[from] StoreError),
91 #[error("pty: {0}")]
92 Pty(String),
93 #[error("run not found: {0}")]
94 NotFound(String),
95 #[error("io: {0}")]
96 Io(#[from] std::io::Error),
97}
98
99#[derive(Debug, Clone)]
103pub struct SpawnOpts {
104 pub cwd: PathBuf,
105 pub env: Vec<(String, String)>,
107 pub label: Option<String>,
108 pub initiator: Initiator,
109 pub pty_cols: u16,
111 pub pty_rows: u16,
113 pub stdin_enabled: bool,
115 pub pin: bool,
117 pub beholder_select: BeholderSelect,
119 pub tty_attached: bool,
122 pub log_fd_enabled: bool,
126 pub origin: Option<String>,
129}
130
131impl Default for SpawnOpts {
132 fn default() -> Self {
133 Self {
134 cwd: std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/")),
135 env: vec![],
136 label: None,
137 initiator: Initiator::Human { camp: "local".to_string() },
138 pty_cols: 80,
139 pty_rows: 24,
140 stdin_enabled: false,
141 pin: false,
142 beholder_select: BeholderSelect::Auto,
143 tty_attached: false,
144 log_fd_enabled: true,
145 origin: None,
146 }
147 }
148}
149
150#[derive(Default)]
155pub struct DriverChannels {
156 pub completion: Option<mpsc::UnboundedSender<(TaskRunId, RunStatus)>>,
159 pub output: Option<mpsc::UnboundedSender<OutputChunk>>,
166}
167
168struct RunControl {
171 kill_tx: mpsc::Sender<KillRequest>,
172 stdin_tx: Option<mpsc::Sender<Vec<u8>>>,
173 master: Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>>,
177}
178
179#[derive(Debug)]
180struct KillRequest {
181 signal: i32,
182}
183
184#[cfg(unix)]
191#[derive(serde::Deserialize)]
192struct ShimRecord {
193 level: String,
194 target: String,
195 msg: String,
196 #[serde(default)]
197 fields: serde_json::Value,
198 #[serde(rename = "_lib", default)]
201 lib: Option<String>,
202 #[serde(rename = "_lib_ver", default)]
204 lib_version: Option<String>,
205}
206
207#[cfg(unix)]
214struct FdCloser(libc::c_int);
215
216#[cfg(unix)]
217impl Drop for FdCloser {
218 fn drop(&mut self) {
219 unsafe { libc::close(self.0) };
220 }
221}
222
223#[cfg(unix)]
226unsafe impl Send for FdCloser {}
227
228pub struct TaskDriver {
234 store: Arc<TaskStore>,
235 active: Arc<Mutex<HashMap<String, RunControl>>>,
236 channels: DriverChannels,
238}
239
240impl TaskDriver {
241 pub async fn new(store: Arc<TaskStore>) -> Result<Self, DriverError> {
246 Self::with_channels(store, DriverChannels::default()).await
247 }
248
249 pub async fn with_channels(
252 store: Arc<TaskStore>,
253 channels: DriverChannels,
254 ) -> Result<Self, DriverError> {
255 let stale = store.list_runs(&RunFilter {
256 status: Some("running".to_string()),
257 ..Default::default()
258 }).await?;
259 for meta in stale {
260 store.update_status(
261 &meta.id,
262 &RunStatus::Lost {
263 reason: "daemon restarted while run was in-flight".to_string(),
264 },
265 ).await?;
266 }
267 Ok(Self {
268 store,
269 active: Arc::new(Mutex::new(HashMap::new())),
270 channels,
271 })
272 }
273
274 pub async fn spawn_run(&self, cmd: &str, opts: SpawnOpts) -> Result<TaskRunId, DriverError> {
287 let id = TaskRunId::new();
288 let started_at = unix_now_secs();
289 let started_at_ms: u64 = started_at.saturating_mul(1000);
290
291 let user_dir = std::env::var_os("YAH_BEHOLDERS_DIR")
294 .map(std::path::PathBuf::from)
295 .or_else(|| {
296 std::env::var_os("HOME")
297 .map(|h| std::path::PathBuf::from(h).join(".yah/beholders"))
298 });
299 let registry = registry_with_user_beholders(user_dir.as_deref());
300 let attach = registry.attach(cmd, &opts.beholder_select, opts.tty_attached);
301 let effective_cmd = if attach.argv.is_empty() {
303 cmd.to_string()
304 } else {
305 attach.argv.join(" ")
306 };
307
308 self.store.insert_run(&TaskRunMeta {
309 id: id.clone(),
310 command: cmd.to_string(),
311 cwd: opts.cwd.clone(),
312 env: opts.env.clone(),
313 started_at,
314 status: RunStatus::Running,
315 label: opts.label.clone(),
316 initiator: opts.initiator.clone(),
317 beholder_status: Some(attach.status),
318 pinned: opts.pin,
319 origin: opts.origin.clone(),
320 }).await?;
321
322 let pty_sys = native_pty_system();
324 let pair = pty_sys
325 .openpty(PtySize {
326 rows: opts.pty_rows,
327 cols: opts.pty_cols,
328 pixel_width: 0,
329 pixel_height: 0,
330 })
331 .map_err(|e| DriverError::Pty(e.to_string()))?;
332
333 let pty_reader = pair
335 .master
336 .try_clone_reader()
337 .map_err(|e| DriverError::Pty(e.to_string()))?;
338
339 let stdin_tx: Option<mpsc::Sender<Vec<u8>>> = if opts.stdin_enabled {
341 let mut writer = pair
342 .master
343 .take_writer()
344 .map_err(|e| DriverError::Pty(e.to_string()))?;
345 let (tx, mut rx) = mpsc::channel::<Vec<u8>>(64);
346 task::spawn(async move {
347 use std::io::Write;
348 while let Some(bytes) = rx.recv().await {
349 let _ = writer.write_all(&bytes);
350 let _ = writer.flush();
351 }
352 });
353 Some(tx)
354 } else {
355 None
356 };
357
358 #[cfg(unix)]
371 let log_fifo: Option<(libc::c_int, FdCloser, std::path::PathBuf)> = if opts.log_fd_enabled {
372 let fifo_path = std::env::temp_dir().join(format!("yah-log-{}.fifo", id));
373 let path_cstr = match std::ffi::CString::new(fifo_path.to_string_lossy().as_bytes()) {
374 Ok(s) => s,
375 Err(_) => {
376 return Err(DriverError::Io(std::io::Error::new(
378 std::io::ErrorKind::InvalidInput,
379 "log FIFO path contained nul byte",
380 )));
381 }
382 };
383 let mkfifo_ret = unsafe { libc::mkfifo(path_cstr.as_ptr(), 0o600) };
384 if mkfifo_ret != 0 {
385 None } else {
387 let rfd = unsafe {
389 libc::open(path_cstr.as_ptr(), libc::O_RDONLY | libc::O_NONBLOCK)
390 };
391 if rfd < 0 {
392 let _ = unsafe { libc::unlink(path_cstr.as_ptr()) };
393 None
394 } else {
395 unsafe { libc::fcntl(rfd, libc::F_SETFL, 0) };
397 let wfd = unsafe {
399 libc::open(path_cstr.as_ptr(), libc::O_WRONLY)
400 };
401 if wfd < 0 {
402 unsafe { libc::close(rfd) };
403 let _ = unsafe { libc::unlink(path_cstr.as_ptr()) };
404 None
405 } else {
406 Some((rfd, FdCloser(wfd), fifo_path))
407 }
408 }
409 }
410 } else {
411 None
412 };
413
414 let mut cb = CommandBuilder::new("sh");
416 cb.args(["-c", &effective_cmd]);
417 cb.cwd(&opts.cwd);
418 for (k, v) in &opts.env {
419 cb.env(k, v);
420 }
421 cb.env("TERM", "xterm-256color");
422
423 #[cfg(unix)]
425 if let Some((_, _, ref fifo_path)) = log_fifo {
426 cb.env("YAH_TASK_RUN", id.to_string());
427 cb.env("YAH_LOG_PIPE", fifo_path.to_string_lossy().as_ref());
428 }
429
430 let child = pair
431 .slave
432 .spawn_command(cb)
433 .map_err(|e| DriverError::Pty(e.to_string()))?;
434 drop(pair.slave);
436
437 let master: Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>> =
440 Arc::new(Mutex::new(pair.master));
441
442 let pid = child.process_id().unwrap_or(0);
443
444 #[cfg(unix)]
452 let log_wfd_holder: Option<FdCloser> = if let Some((rfd, wfd, fifo_path)) = log_fifo {
453 let store_log = Arc::clone(&self.store);
454 let id_log = id.clone();
455 let rt = tokio::runtime::Handle::current();
456 tokio::task::spawn_blocking(move || {
459 run_log_receiver(rt, store_log, id_log, rfd, fifo_path, started_at_ms);
460 });
461 Some(wfd)
462 } else {
463 None
464 };
465
466 let (kill_tx, kill_rx) = mpsc::channel::<KillRequest>(4);
468 let (reader_done_tx, reader_done_rx) = oneshot::channel::<()>();
469
470 {
473 let store_r = Arc::clone(&self.store);
474 let id_r = id.clone();
475 let mut beholder = attach.beholder;
476 let output_tx = self.channels.output.clone();
477 let rt = tokio::runtime::Handle::current();
478 tokio::task::spawn_blocking(move || {
479 let mut buf = [0u8; READ_BUF_SIZE];
480 let mut reader = pty_reader;
481 loop {
482 match reader.read(&mut buf) {
483 Ok(0) | Err(_) => break,
484 Ok(n) => {
485 let offset = elapsed_ms(started_at_ms);
486 let append_res = rt.block_on(store_r.append_chunk(
487 &id_r,
488 offset,
489 Stream::Stdout,
490 &buf[..n],
491 ));
492 if let Ok(seq) = append_res {
493 let chunk = (output_tx.is_some() || beholder.is_some()).then(|| {
497 OutputChunk {
498 run_id: id_r.clone(),
499 seq,
500 offset_ms: offset,
501 stream: Stream::Stdout,
502 bytes: buf[..n].to_vec(),
503 }
504 });
505 if let (Some(tx), Some(c)) = (&output_tx, &chunk) {
509 let _ = tx.send(c.clone());
510 }
511 let mut detach_beholder = false;
512 if let (Some(b), Some(chunk)) = (beholder.as_mut(), &chunk) {
513 for ev in b.parse_chunk(chunk) {
514 let _ = rt.block_on(store_r.append_event(
515 &ev.run_id,
516 ev.offset_ms,
517 ev.level,
518 &ev.target,
519 &ev.msg,
520 &ev.fields,
521 ev.anchor.as_ref().map(|a| a.seq),
522 &ev.source,
523 ));
524 }
525 if let Some(reason) = b.unknown_format_reason() {
526 let new_status = BeholderStatus::unknown_format_with_reason(
527 b.name(),
528 reason,
529 );
530 let _ = rt.block_on(
531 store_r.update_beholder_status(&id_r, &new_status),
532 );
533 detach_beholder = true;
534 }
535 }
536 if detach_beholder {
537 beholder = None;
538 }
539 }
540 }
541 }
542 }
543 if let Some(ref mut b) = beholder {
544 let final_offset = elapsed_ms(started_at_ms);
545 for ev in b.on_done(&id_r, final_offset) {
546 let _ = rt.block_on(store_r.append_event(
547 &ev.run_id,
548 ev.offset_ms,
549 ev.level,
550 &ev.target,
551 &ev.msg,
552 &ev.fields,
553 ev.anchor.as_ref().map(|a| a.seq),
554 &ev.source,
555 ));
556 }
557 if let Some(reason) = b.unknown_format_reason() {
558 let new_status = BeholderStatus::unknown_format_with_reason(b.name(), reason);
559 let _ = rt.block_on(store_r.update_beholder_status(&id_r, &new_status));
560 }
561 }
562 let _ = reader_done_tx.send(());
563 });
564 }
565
566 {
570 let store_l = Arc::clone(&self.store);
571 let active_l = Arc::clone(&self.active);
572 let id_l = id.clone();
573 let master_l = Arc::clone(&master);
574 let completion_tx_l = self.channels.completion.clone();
575 #[cfg(unix)]
576 let wfd_l = log_wfd_holder;
577 task::spawn(async move {
578 run_lifecycle(
579 store_l,
580 active_l,
581 id_l,
582 pid,
583 child,
584 master_l,
585 kill_rx,
586 reader_done_rx,
587 completion_tx_l,
588 #[cfg(unix)]
589 wfd_l,
590 )
591 .await;
592 });
593 }
594
595 self.active
596 .lock()
597 .unwrap()
598 .insert(id.to_string(), RunControl { kill_tx, stdin_tx, master });
599
600 Ok(id)
601 }
602
603 pub async fn resize_run(
610 &self,
611 id: &TaskRunId,
612 cols: u16,
613 rows: u16,
614 ) -> Result<(), DriverError> {
615 let master = self
616 .active
617 .lock()
618 .unwrap()
619 .get(&id.to_string())
620 .map(|c| Arc::clone(&c.master));
621
622 match master {
623 Some(m) => {
624 let size = PtySize { rows, cols, pixel_width: 0, pixel_height: 0 };
625 m.lock()
626 .unwrap()
627 .resize(size)
628 .map_err(|e| DriverError::Pty(e.to_string()))
629 }
630 None => Err(DriverError::NotFound(id.to_string())),
631 }
632 }
633
634 pub async fn kill_run(&self, id: &TaskRunId, signal: Option<i32>) -> Result<(), DriverError> {
641 let kill_tx = self
642 .active
643 .lock()
644 .unwrap()
645 .get(&id.to_string())
646 .map(|c| c.kill_tx.clone());
647
648 match kill_tx {
649 Some(tx) => tx
650 .send(KillRequest { signal: signal.unwrap_or(SIGTERM) })
651 .await
652 .map_err(|_| DriverError::NotFound(id.to_string())),
653 None => Err(DriverError::NotFound(id.to_string())),
654 }
655 }
656
657 pub async fn send_stdin(&self, id: &TaskRunId, bytes: Vec<u8>) -> Result<(), DriverError> {
659 let stdin_tx = self
660 .active
661 .lock()
662 .unwrap()
663 .get(&id.to_string())
664 .and_then(|c| c.stdin_tx.clone());
665
666 match stdin_tx {
667 Some(tx) => tx
668 .send(bytes)
669 .await
670 .map_err(|_| DriverError::NotFound(id.to_string())),
671 None => Err(DriverError::NotFound(id.to_string())),
672 }
673 }
674}
675
676#[cfg(unix)]
685fn run_log_receiver(
686 rt: tokio::runtime::Handle,
687 store: Arc<TaskStore>,
688 run_id: TaskRunId,
689 read_fd: libc::c_int,
690 fifo_path: std::path::PathBuf,
691 started_at_ms: u64,
692) {
693 use std::io::BufRead;
694 use std::os::unix::io::FromRawFd;
695
696 let file = unsafe { std::fs::File::from_raw_fd(read_fd) };
699 let reader = std::io::BufReader::new(file);
700
701 for line in reader.lines() {
702 let line = match line {
703 Ok(l) => l,
704 Err(_) => break,
705 };
706 let trimmed = line.trim();
707 if trimmed.is_empty() {
708 continue;
709 }
710 let rec: ShimRecord = match serde_json::from_str(trimmed) {
711 Ok(r) => r,
712 Err(_) => continue, };
714 let level = rec.level.parse::<crate::types::Level>().unwrap_or(crate::types::Level::Info);
715 let source = crate::types::EventSource::Shim {
716 lib: rec.lib.unwrap_or_else(|| "unknown".to_string()),
717 version: rec.lib_version.unwrap_or_else(|| "0.0.0".to_string()),
718 };
719 let fields = if rec.fields.is_object() {
720 rec.fields
721 } else {
722 serde_json::Value::Object(Default::default())
723 };
724 let offset = elapsed_ms(started_at_ms);
725 let _ = rt.block_on(store.append_event(
726 &run_id,
727 offset,
728 level,
729 &rec.target,
730 &rec.msg,
731 &fields,
732 None,
733 &source,
734 ));
735 }
736
737 let _ = std::fs::remove_file(&fifo_path);
739}
740
741async fn run_lifecycle(
744 store: Arc<TaskStore>,
745 active: Arc<Mutex<HashMap<String, RunControl>>>,
746 id: TaskRunId,
747 pid: u32,
748 child: Box<dyn portable_pty::Child + Send>,
749 master: Arc<Mutex<Box<dyn portable_pty::MasterPty + Send>>>,
750 mut kill_rx: mpsc::Receiver<KillRequest>,
751 reader_done_rx: oneshot::Receiver<()>,
752 completion_tx: Option<tokio::sync::mpsc::UnboundedSender<(TaskRunId, RunStatus)>>,
753 #[cfg(unix)]
757 _log_wfd: Option<FdCloser>,
758) {
759 let reader_done = async { reader_done_rx.await.ok(); };
762 tokio::pin!(reader_done);
763
764 let sent_signal: Option<i32>;
765
766 tokio::select! {
767 req = kill_rx.recv() => {
768 match req {
769 Some(KillRequest { signal }) => {
770 send_unix_signal(pid, signal);
771 if signal == SIGKILL {
772 sent_signal = Some(SIGKILL);
773 } else {
774 tokio::select! {
776 _ = &mut reader_done => {
777 sent_signal = Some(signal);
779 }
780 _ = tokio::time::sleep(DEFAULT_GRACE) => {
781 send_unix_signal(pid, SIGKILL);
783 sent_signal = Some(SIGKILL);
784 }
785 }
786 }
787 }
788 None => {
790 send_unix_signal(pid, SIGKILL);
791 sent_signal = Some(SIGKILL);
792 }
793 }
794 }
795 _ = &mut reader_done => {
796 sent_signal = None;
797 }
798 }
799
800 let exit_code = task::spawn_blocking(move || {
805 let mut c = child;
806 let _m = master; c.wait().ok().map(|s| s.exit_code())
808 })
809 .await
810 .ok()
811 .flatten();
812
813 let ended_at = unix_now_secs();
814 let status = match sent_signal {
815 Some(sig) => RunStatus::Killed { signal: sig, ended_at },
816 None => match exit_code {
817 Some(code) => RunStatus::Done { exit_code: code as i32, ended_at },
818 None => RunStatus::Lost {
819 reason: "process exited without an exit code".to_string(),
820 },
821 },
822 };
823
824 if let Err(e) = store.update_status(&id, &status).await {
830 eprintln!("[yah task-runs] failed to record terminal status for run {id}: {e}");
831 }
832 if let Some(ref tx) = completion_tx {
833 let _ = tx.send((id.clone(), status));
834 }
835 active.lock().unwrap().remove(&id.to_string());
836}
837
838fn send_unix_signal(pid: u32, signal: i32) {
841 #[cfg(unix)]
842 unsafe {
843 libc::kill(pid as libc::pid_t, signal);
844 }
845 }
847
848fn unix_now_secs() -> u64 {
849 SystemTime::now()
850 .duration_since(UNIX_EPOCH)
851 .unwrap_or_default()
852 .as_secs()
853}
854
855fn elapsed_ms(started_at_ms: u64) -> u32 {
856 let now_ms = SystemTime::now()
857 .duration_since(UNIX_EPOCH)
858 .unwrap_or_default()
859 .as_millis() as u64;
860 now_ms.saturating_sub(started_at_ms).min(u32::MAX as u64) as u32
861}
862
863#[cfg(test)]
866mod tests {
867 use super::*;
868 use crate::store::ChunkFilter;
869
870 async fn open_store(dir: &tempfile::TempDir) -> Arc<TaskStore> {
871 Arc::new(TaskStore::open(&dir.path().join("tr.turso")).await.unwrap())
872 }
873
874 #[tokio::test]
877 async fn lost_on_disappear_marks_stale_running_runs() {
878 let dir = tempfile::tempdir().unwrap();
879 let store = open_store(&dir).await;
880
881 let stale_id = TaskRunId::new();
883 store
884 .insert_run(&TaskRunMeta {
885 id: stale_id.clone(),
886 command: "sleep 9999".to_string(),
887 cwd: "/tmp".into(),
888 env: vec![],
889 started_at: unix_now_secs() - 60,
890 status: RunStatus::Running,
891 label: None,
892 initiator: Initiator::Human { camp: "test".to_string() },
893 beholder_status: None,
894 pinned: false,
895 origin: None,
896 })
897 .await
898 .unwrap();
899
900 let _driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
902
903 let meta = store.get_run(&stale_id).await.unwrap().unwrap();
904 assert!(
905 matches!(meta.status, RunStatus::Lost { .. }),
906 "stale run should be Lost, got {:?}",
907 meta.status
908 );
909 }
910
911 #[tokio::test]
912 async fn new_driver_does_not_touch_completed_runs() {
913 let dir = tempfile::tempdir().unwrap();
914 let store = open_store(&dir).await;
915
916 let done_id = TaskRunId::new();
917 store
918 .insert_run(&TaskRunMeta {
919 id: done_id.clone(),
920 command: "true".to_string(),
921 cwd: "/tmp".into(),
922 env: vec![],
923 started_at: unix_now_secs() - 10,
924 status: RunStatus::Running,
925 label: None,
926 initiator: Initiator::Human { camp: "test".to_string() },
927 beholder_status: None,
928 pinned: false,
929 origin: None,
930 })
931 .await
932 .unwrap();
933 store
934 .update_status(&done_id, &RunStatus::Done { exit_code: 0, ended_at: unix_now_secs() })
935 .await
936 .unwrap();
937
938 let _driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
939
940 let meta = store.get_run(&done_id).await.unwrap().unwrap();
941 assert!(
942 matches!(meta.status, RunStatus::Done { .. }),
943 "completed run must not be touched"
944 );
945 }
946
947 #[tokio::test]
950 async fn spawn_echo_and_read_chunks() {
951 let dir = tempfile::tempdir().unwrap();
952 let store = open_store(&dir).await;
953 let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
954
955 let id = driver
956 .spawn_run(
957 "echo hello_world",
958 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
959 )
960 .await
961 .unwrap();
962
963 let deadline = std::time::Instant::now() + Duration::from_secs(5);
965 loop {
966 let meta = store.get_run(&id).await.unwrap().unwrap();
967 if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
968 break;
969 }
970 if std::time::Instant::now() > deadline {
971 panic!("run did not complete in time, status={:?}", meta.status);
972 }
973 tokio::time::sleep(Duration::from_millis(50)).await;
974 }
975
976 let chunks = store
978 .get_chunks(&id, &ChunkFilter::default())
979 .await
980 .unwrap();
981 let output: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
982 let text = String::from_utf8_lossy(&output);
983 assert!(
984 text.contains("hello_world"),
985 "expected 'hello_world' in output, got: {text:?}"
986 );
987
988 let meta = store.get_run(&id).await.unwrap().unwrap();
989 assert!(
990 matches!(meta.status, RunStatus::Done { exit_code: 0, .. }),
991 "expected Done(0), got {:?}",
992 meta.status
993 );
994 }
995
996 #[tokio::test]
997 async fn spawn_failing_command_records_nonzero_exit() {
998 let dir = tempfile::tempdir().unwrap();
999 let store = open_store(&dir).await;
1000 let driver = TaskDriver::new(Arc::clone(&store)).await.unwrap();
1001
1002 let id = driver
1003 .spawn_run(
1004 "exit 42",
1005 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
1006 )
1007 .await
1008 .unwrap();
1009
1010 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1011 loop {
1012 let meta = store.get_run(&id).await.unwrap().unwrap();
1013 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1014 match meta.status {
1015 RunStatus::Done { exit_code, .. } => {
1016 assert_ne!(exit_code, 0, "exit 42 should produce a non-zero exit code");
1017 }
1018 other => panic!("unexpected status: {other:?}"),
1019 }
1020 break;
1021 }
1022 if std::time::Instant::now() > deadline {
1023 panic!("run did not complete in time");
1024 }
1025 tokio::time::sleep(Duration::from_millis(50)).await;
1026 }
1027 }
1028
1029 #[cfg(unix)]
1032 #[tokio::test]
1033 async fn kill_with_sigterm_transitions_to_killed() {
1034 let dir = tempfile::tempdir().unwrap();
1035 let store = open_store(&dir).await;
1036 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1037
1038 let id = driver
1039 .spawn_run(
1040 "sleep 60",
1041 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
1042 )
1043 .await
1044 .unwrap();
1045
1046 tokio::time::sleep(Duration::from_millis(100)).await;
1048
1049 driver.kill_run(&id, Some(SIGTERM)).await.unwrap();
1050
1051 let deadline = std::time::Instant::now() + Duration::from_secs(10);
1052 loop {
1053 let meta = store.get_run(&id).await.unwrap().unwrap();
1054 if matches!(meta.status, RunStatus::Killed { .. } | RunStatus::Lost { .. }) {
1055 assert!(
1056 matches!(meta.status, RunStatus::Killed { .. }),
1057 "expected Killed, got {:?}",
1058 meta.status
1059 );
1060 break;
1061 }
1062 if std::time::Instant::now() > deadline {
1063 panic!("run did not become Killed in time, status={:?}", meta.status);
1064 }
1065 tokio::time::sleep(Duration::from_millis(50)).await;
1066 }
1067 }
1068
1069 #[cfg(unix)]
1070 #[tokio::test]
1071 async fn kill_run_returns_not_found_after_exit() {
1072 let dir = tempfile::tempdir().unwrap();
1073 let store = open_store(&dir).await;
1074 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1075
1076 let id = driver
1077 .spawn_run(
1078 "echo done",
1079 SpawnOpts { cwd: "/tmp".into(), ..Default::default() },
1080 )
1081 .await
1082 .unwrap();
1083
1084 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1086 loop {
1087 let meta = store.get_run(&id).await.unwrap().unwrap();
1088 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1089 break;
1090 }
1091 if std::time::Instant::now() > deadline {
1092 panic!("run did not complete");
1093 }
1094 tokio::time::sleep(Duration::from_millis(50)).await;
1095 }
1096
1097 let result = driver.kill_run(&id, None).await;
1099 assert!(
1100 matches!(result, Err(DriverError::NotFound(_))),
1101 "expected NotFound, got {result:?}"
1102 );
1103 }
1104
1105 #[cfg(unix)]
1108 #[tokio::test]
1109 async fn stdin_send_reaches_child() {
1110 let dir = tempfile::tempdir().unwrap();
1111 let store = open_store(&dir).await;
1112 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1113
1114 let id = driver
1116 .spawn_run(
1117 "read line && echo got_$line",
1118 SpawnOpts {
1119 cwd: "/tmp".into(),
1120 stdin_enabled: true,
1121 ..Default::default()
1122 },
1123 )
1124 .await
1125 .unwrap();
1126
1127 tokio::time::sleep(Duration::from_millis(150)).await;
1128 driver.send_stdin(&id, b"hello\n".to_vec()).await.unwrap();
1129
1130 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1131 loop {
1132 let meta = store.get_run(&id).await.unwrap().unwrap();
1133 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1134 break;
1135 }
1136 if std::time::Instant::now() > deadline {
1137 panic!("run did not complete after stdin input");
1138 }
1139 tokio::time::sleep(Duration::from_millis(50)).await;
1140 }
1141
1142 let chunks = store.get_chunks(&id, &ChunkFilter::default()).await.unwrap();
1143 let raw: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
1144 let text = String::from_utf8_lossy(&raw);
1145 assert!(
1146 text.contains("got_hello"),
1147 "expected 'got_hello' in output, got: {text:?}"
1148 );
1149 }
1150
1151 #[tokio::test]
1155 async fn resize_run_changes_geometry_the_child_sees() {
1156 let dir = tempfile::tempdir().unwrap();
1157 let store = open_store(&dir).await;
1158 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1159
1160 let id = driver
1162 .spawn_run(
1163 "read line && stty size",
1164 SpawnOpts {
1165 cwd: "/tmp".into(),
1166 stdin_enabled: true,
1167 ..Default::default()
1170 },
1171 )
1172 .await
1173 .unwrap();
1174
1175 tokio::time::sleep(Duration::from_millis(150)).await;
1176 driver.resize_run(&id, 120, 40).await.unwrap();
1177 driver.send_stdin(&id, b"go\n".to_vec()).await.unwrap();
1178
1179 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1180 loop {
1181 let meta = store.get_run(&id).await.unwrap().unwrap();
1182 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1183 break;
1184 }
1185 if std::time::Instant::now() > deadline {
1186 panic!("run did not complete after stdin input");
1187 }
1188 tokio::time::sleep(Duration::from_millis(50)).await;
1189 }
1190
1191 let chunks = store.get_chunks(&id, &ChunkFilter::default()).await.unwrap();
1192 let raw: Vec<u8> = chunks.into_iter().flat_map(|c| c.bytes).collect();
1193 let text = String::from_utf8_lossy(&raw);
1194 assert!(
1195 text.contains("40 120"),
1196 "expected resized geometry '40 120' in output, got: {text:?}"
1197 );
1198 }
1199
1200 #[tokio::test]
1203 async fn resize_run_returns_not_found_after_exit() {
1204 let dir = tempfile::tempdir().unwrap();
1205 let store = open_store(&dir).await;
1206 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1207
1208 let id = driver
1209 .spawn_run("true", SpawnOpts { cwd: "/tmp".into(), ..Default::default() })
1210 .await
1211 .unwrap();
1212
1213 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1214 loop {
1215 let meta = store.get_run(&id).await.unwrap().unwrap();
1216 if !matches!(meta.status, RunStatus::Running | RunStatus::Pending) {
1217 break;
1218 }
1219 if std::time::Instant::now() > deadline {
1220 panic!("run did not exit");
1221 }
1222 tokio::time::sleep(Duration::from_millis(50)).await;
1223 }
1224
1225 assert!(matches!(
1226 driver.resize_run(&id, 100, 30).await,
1227 Err(DriverError::NotFound(_))
1228 ));
1229 }
1230
1231 #[cfg(unix)]
1239 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
1240 async fn log_pipe_events_land_in_store() {
1241 use crate::store::EventFilter;
1242
1243 let dir = tempfile::tempdir().unwrap();
1244 let store = open_store(&dir).await;
1245 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1246
1247 let cmd = r#"printf '{"level":"warn","target":"test.shim","msg":"hello-from-pipe","fields":{"x":42},"_lib":"test-shim","_lib_ver":"0.1.0"}\n' >> "$YAH_LOG_PIPE""#;
1250
1251 let id = driver
1252 .spawn_run(cmd, SpawnOpts { cwd: "/tmp".into(), ..Default::default() })
1253 .await
1254 .unwrap();
1255
1256 let deadline = std::time::Instant::now() + Duration::from_secs(20);
1260 loop {
1261 let meta = store.get_run(&id).await.unwrap().unwrap();
1262 if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
1263 break;
1264 }
1265 if std::time::Instant::now() > deadline {
1266 panic!("run did not complete in time");
1267 }
1268 tokio::time::sleep(Duration::from_millis(50)).await;
1269 }
1270
1271 tokio::time::sleep(Duration::from_millis(500)).await;
1274
1275 let events = store.query_events(&id, &EventFilter::default()).await.unwrap();
1276 assert!(
1277 !events.is_empty(),
1278 "expected at least one shim event, got none"
1279 );
1280 let ev = events.iter().find(|e| e.target == "test.shim");
1281 let ev = ev.expect("event with target 'test.shim' not found");
1282 assert_eq!(ev.msg, "hello-from-pipe");
1283 assert_eq!(ev.level, crate::types::Level::Warn);
1284 assert!(
1285 matches!(&ev.source, crate::types::EventSource::Shim { lib, .. } if lib == "test-shim"),
1286 "unexpected source: {:?}",
1287 ev.source
1288 );
1289 assert_eq!(ev.fields.get("x"), Some(&serde_json::json!(42)));
1290 }
1291
1292 #[cfg(unix)]
1295 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1296 async fn log_pipe_disabled_produces_no_events() {
1297 use crate::store::EventFilter;
1298
1299 let dir = tempfile::tempdir().unwrap();
1300 let store = open_store(&dir).await;
1301 let driver = Arc::new(TaskDriver::new(Arc::clone(&store)).await.unwrap());
1302
1303 let cmd = r#"[ -n "$YAH_LOG_PIPE" ] && printf '{"level":"info","target":"t","msg":"m","fields":{}}\n' >> "$YAH_LOG_PIPE" || true"#;
1306
1307 let id = driver
1308 .spawn_run(
1309 cmd,
1310 SpawnOpts { cwd: "/tmp".into(), log_fd_enabled: false, ..Default::default() },
1311 )
1312 .await
1313 .unwrap();
1314
1315 let deadline = std::time::Instant::now() + Duration::from_secs(5);
1316 loop {
1317 let meta = store.get_run(&id).await.unwrap().unwrap();
1318 if matches!(meta.status, RunStatus::Done { .. } | RunStatus::Lost { .. }) {
1319 break;
1320 }
1321 if std::time::Instant::now() > deadline {
1322 panic!("run did not complete");
1323 }
1324 tokio::time::sleep(Duration::from_millis(50)).await;
1325 }
1326
1327 tokio::time::sleep(Duration::from_millis(100)).await;
1328
1329 let events = store.query_events(&id, &EventFilter::default()).await.unwrap();
1330 assert!(
1331 events.is_empty(),
1332 "expected no shim events when log_fd_enabled=false, got {}",
1333 events.len()
1334 );
1335 }
1336}