IdeaSDK/crates/infrastructure/tests/background_task_runner.rs

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<i64>);

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<PtyHandle>,
    status: Option<ExitStatus>,
    kills: usize,
    scrollback: Vec<u8>,
}

#[derive(Default)]
struct FakePty {
    state: Mutex<FakePtyState>,
    wait_changed: tokio::sync::Notify,
    subscribe_count: AtomicUsize,
}

impl FakePty {
    fn complete(&self, code: Option<i32>) {
        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 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<PtyHandle, PtyError> {
        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<OutputStream, PtyError> {
        self.subscribe_count.fetch_add(1, Ordering::SeqCst);
        Ok(Box::new(std::iter::empty()))
    }

    fn scrollback(&self, _handle: &PtyHandle) -> Result<Vec<u8>, PtyError> {
        Ok(self.state.lock().unwrap().scrollback.clone())
    }

    async fn wait(&self, _handle: &PtyHandle) -> Result<ExitStatus, PtyError> {
        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<Option<ExitStatus>, PtyError> {
        Ok(self.state.lock().unwrap().status)
    }

    async fn kill(&self, _handle: &PtyHandle) -> Result<ExitStatus, PtyError> {
        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<u64>) -> 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<FakePty>,
    clock: Arc<FakeClock>,
) -> (
    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 = runner
        .pty_handle_for(task_id)
        .expect("runner exposes live PTY 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 double_ui_subscriber_uses_independent_pty_streams() {
    let pty = Arc::new(FakePty::default());
    let clock = Arc::new(FakeClock::default());
    let (runner, _completions) = runner_with(Arc::clone(&pty), clock);
    let task_id = TaskId::from_uuid(id(14));

    runner
        .spawn(spawn_spec(task_id, None))
        .await
        .expect("spawn succeeds");
    let handle = runner
        .pty_handle_for(task_id)
        .expect("runner exposes live PTY handle");

    let _first = pty.subscribe_output(&handle).expect("first ui subscribes");
    let _second = pty.subscribe_output(&handle).expect("second ui subscribes");

    assert_eq!(
        pty.subscribe_count(),
        2,
        "each UI attach must get its own PTY subscription"
    );
    assert_eq!(
        runner.pty_handle_for(task_id),
        Some(handle),
        "UI subscribers must not remove the runner's live handle"
    );
}

#[tokio::test]
async fn completed_task_no_longer_exposes_live_pty_handle_for_attach() {
    let pty = Arc::new(FakePty::default());
    pty.set_scrollback(b"final tail");
    let clock = Arc::new(FakeClock::default());
    let (runner, mut completions) = runner_with(Arc::clone(&pty), clock);
    let task_id = TaskId::from_uuid(id(15));

    runner
        .spawn(spawn_spec(task_id, None))
        .await
        .expect("spawn succeeds");
    assert!(runner.pty_handle_for(task_id).is_some());
    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_eq!(runner.pty_handle_for(task_id), None);
    assert_eq!(
        pty.subscribe_count(),
        0,
        "attach after completion must not subscribe to a dead PTY handle"
    );
}

#[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);
}