use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; use async_trait::async_trait; use domain::ports::{ BackgroundTaskRunner, Clock, ExitStatus, OutputStream, PtyError, PtyHandle, PtyPort, SpawnSpec, }; use domain::{ AgentId, BackgroundTaskKind, BackgroundTaskResult, BackgroundTaskWakePolicy, ProjectId, ProjectPath, PtySize, SessionId, TaskId, }; use infrastructure::CommandBackgroundRunner; use uuid::Uuid; #[derive(Default)] struct FakeClock(Mutex); impl FakeClock { fn set(&self, now: i64) { *self.0.lock().unwrap() = now; } } impl Clock for FakeClock { fn now_millis(&self) -> i64 { *self.0.lock().unwrap() } } #[derive(Default)] struct FakePtyState { handle: Option, status: Option, kills: usize, scrollback: Vec, } #[derive(Default)] struct FakePty { state: Mutex, wait_changed: tokio::sync::Notify, subscribe_count: AtomicUsize, } impl FakePty { fn complete(&self, code: Option) { self.state.lock().unwrap().status = Some(ExitStatus { code }); self.wait_changed.notify_waiters(); } fn set_scrollback(&self, bytes: &[u8]) { self.state.lock().unwrap().scrollback = bytes.to_vec(); } fn last_handle(&self) -> PtyHandle { self.state .lock() .unwrap() .handle .clone() .expect("spawned handle") } fn kill_count(&self) -> usize { self.state.lock().unwrap().kills } fn subscribe_count(&self) -> usize { self.subscribe_count.load(Ordering::SeqCst) } } #[async_trait] impl PtyPort for FakePty { async fn spawn(&self, _spec: SpawnSpec, _size: PtySize) -> Result { let handle = PtyHandle { session_id: SessionId::from_uuid(Uuid::from_u128(0x777)), }; self.state.lock().unwrap().handle = Some(handle.clone()); Ok(handle) } fn write(&self, _handle: &PtyHandle, _data: &[u8]) -> Result<(), PtyError> { Ok(()) } fn resize(&self, _handle: &PtyHandle, _size: PtySize) -> Result<(), PtyError> { Ok(()) } fn subscribe_output(&self, _handle: &PtyHandle) -> Result { self.subscribe_count.fetch_add(1, Ordering::SeqCst); Ok(Box::new(std::iter::empty())) } fn scrollback(&self, _handle: &PtyHandle) -> Result, PtyError> { Ok(self.state.lock().unwrap().scrollback.clone()) } async fn wait(&self, _handle: &PtyHandle) -> Result { loop { if let Some(status) = self.state.lock().unwrap().status { return Ok(status); } self.wait_changed.notified().await; } } fn try_wait(&self, _handle: &PtyHandle) -> Result, PtyError> { Ok(self.state.lock().unwrap().status) } async fn kill(&self, _handle: &PtyHandle) -> Result { let mut state = self.state.lock().unwrap(); state.kills += 1; let status = state.status.get_or_insert(ExitStatus { code: Some(130) }); Ok(*status) } } fn id(n: u128) -> Uuid { Uuid::from_u128(n) } fn spawn_spec(task_id: TaskId, deadline_ms: Option) -> domain::ports::BackgroundTaskSpec { domain::ports::BackgroundTaskSpec { task_id, project_id: ProjectId::from_uuid(id(1)), owner_agent_id: AgentId::from_uuid(id(2)), kind: BackgroundTaskKind::Command { label: "test command".to_owned(), }, wake_policy: BackgroundTaskWakePolicy::RecordOnly, command: Some(SpawnSpec { command: "/bin/sh".to_owned(), args: vec!["-c".to_owned(), "echo test".to_owned()], cwd: ProjectPath::new("/tmp").unwrap(), env: Vec::new(), context_plan: None, sandbox: None, }), deadline_ms, } } fn runner_with( pty: Arc, clock: Arc, ) -> ( CommandBackgroundRunner, domain::ports::BackgroundCompletionStream, ) { let runner = CommandBackgroundRunner::new(pty, clock); let completions = runner.subscribe_completions(); (runner, completions) } #[tokio::test] async fn completes_from_wait_without_draining_output() { let pty = Arc::new(FakePty::default()); pty.set_scrollback(b"hello from scrollback"); let clock = Arc::new(FakeClock::default()); let (runner, mut completions) = runner_with(Arc::clone(&pty), clock); let task_id = TaskId::from_uuid(id(10)); runner .spawn(spawn_spec(task_id, None)) .await .expect("spawn succeeds"); pty.complete(Some(0)); let completion = tokio::time::timeout( Duration::from_secs(1), tokio::task::spawn_blocking(move || completions.next().expect("completion")), ) .await .expect("completion arrives") .expect("completion thread joins"); assert_eq!(completion.task_id, task_id); match completion.result { BackgroundTaskResult::Success { exit_code, stdout_tail, .. } => { assert_eq!(exit_code, Some(0)); assert_eq!(stdout_tail.as_deref(), Some("hello from scrollback")); } other => panic!("expected success, got {other:?}"), } assert_eq!(pty.subscribe_count(), 0, "runner must not drain output"); assert_eq!(pty.kill_count(), 1, "natural exit is cleaned up once"); } #[tokio::test] async fn ui_subscriber_does_not_interfere_with_runner_completion() { let pty = Arc::new(FakePty::default()); let clock = Arc::new(FakeClock::default()); let (runner, mut completions) = runner_with(Arc::clone(&pty), clock); let task_id = TaskId::from_uuid(id(11)); runner .spawn(spawn_spec(task_id, None)) .await .expect("spawn succeeds"); let handle = pty.last_handle(); let _ui_stream = pty.subscribe_output(&handle).expect("ui subscribes"); pty.complete(Some(0)); let completion = tokio::time::timeout( Duration::from_secs(1), tokio::task::spawn_blocking(move || completions.next().expect("completion")), ) .await .expect("completion arrives") .expect("completion thread joins"); assert_eq!(completion.task_id, task_id); assert!(matches!( completion.result, BackgroundTaskResult::Success { .. } )); assert_eq!( pty.subscribe_count(), 1, "only the explicit UI subscription should exist" ); } #[tokio::test] async fn deadline_expires_and_kills_when_wait_never_resolves() { let pty = Arc::new(FakePty::default()); let clock = Arc::new(FakeClock::default()); clock.set(100); let (runner, mut completions) = runner_with(Arc::clone(&pty), Arc::clone(&clock)); let task_id = TaskId::from_uuid(id(12)); runner .spawn(spawn_spec(task_id, Some(101))) .await .expect("spawn succeeds"); let completion = tokio::time::timeout( Duration::from_secs(1), tokio::task::spawn_blocking(move || completions.next().expect("completion")), ) .await .expect("completion arrives") .expect("completion thread joins"); assert_eq!(completion.task_id, task_id); assert!(matches!( completion.result, BackgroundTaskResult::Expired { .. } )); assert_eq!(pty.kill_count(), 1); assert_eq!(pty.subscribe_count(), 0); } #[tokio::test] async fn cancel_wins_even_if_output_is_still_available() { let pty = Arc::new(FakePty::default()); pty.set_scrollback(b"still streaming"); let clock = Arc::new(FakeClock::default()); let (runner, mut completions) = runner_with(Arc::clone(&pty), clock); let task_id = TaskId::from_uuid(id(13)); runner .spawn(spawn_spec(task_id, None)) .await .expect("spawn succeeds"); runner.cancel(task_id).await.expect("cancel succeeds"); let completion = tokio::time::timeout( Duration::from_secs(1), tokio::task::spawn_blocking(move || completions.next().expect("completion")), ) .await .expect("completion arrives") .expect("completion thread joins"); assert_eq!(completion.task_id, task_id); assert!(matches!( completion.result, BackgroundTaskResult::Cancelled { .. } )); assert_eq!(pty.kill_count(), 1); assert_eq!(pty.subscribe_count(), 0); }