batty-cli 0.11.63

//! Integration tests for the shim runtime.
//!
//! Gated behind `shim-integration` feature flag.
//! Run with: cargo test --features shim-integration
//!
//! These tests spawn a real bash process via the shim runtime and
//! exercise the protocol end-to-end.

use std::fs;
use std::os::unix::io::IntoRawFd;
use std::os::unix::process::CommandExt;
use std::path::PathBuf;
use std::process::{Child, Command as ProcessCommand, Stdio};
use std::time::Duration;

use crate::shim::protocol::{self, Channel, Command, Event, ShimState};
use crate::shim::runtime::ShimArgs;

/// Check if an error is a read timeout (WouldBlock/TimedOut).
fn is_timeout_error(e: &anyhow::Error) -> bool {
    if let Some(io_err) = e.downcast_ref::<std::io::Error>() {
        matches!(
            io_err.kind(),
            std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
        )
    } else {
        false
    }
}

/// Helper: spawn a shim with bash in a background thread, return the parent channel.
/// Sets a 500ms read timeout so recv() doesn't block forever.
fn spawn_bash_shim() -> Channel {
    let (parent_sock, child_sock) = protocol::socketpair().unwrap();
    let channel = Channel::new(child_sock);

    let args = ShimArgs {
        id: "test-agent".into(),
        agent_type: crate::shim::classifier::AgentType::Generic,
        cmd: "bash --norc --noprofile".into(),
        cwd: PathBuf::from("/tmp"),
        rows: 24,
        cols: 80,
        pty_log_path: None,
        graceful_shutdown_timeout_secs: 5,
        auto_commit_on_restart: true,
    };

    std::thread::spawn(move || {
        crate::shim::runtime::run(args, channel).ok();
    });

    let mut ch = Channel::new(parent_sock);
    ch.set_read_timeout(Some(Duration::from_millis(500)))
        .unwrap();
    ch
}

fn build_batty_binary() -> PathBuf {
    let repo_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let status = ProcessCommand::new("cargo")
        .args(["build", "--quiet", "--bin", "batty"])
        .current_dir(&repo_root)
        .status()
        .expect("failed to build batty binary");
    assert!(status.success(), "cargo build --bin batty should succeed");
    repo_root.join("target").join("debug").join("batty")
}

fn spawn_external_shim(cmd: &str, cwd: &std::path::Path) -> (Child, Channel) {
    let (parent_sock, child_sock) = protocol::socketpair().unwrap();
    let child_fd = child_sock.into_raw_fd();
    let batty = build_batty_binary();

    let process = unsafe {
        let child_fd_copy = child_fd;
        ProcessCommand::new(&batty)
            .args([
                "shim",
                "--id",
                "test-agent",
                "--agent-type",
                "generic",
                "--cmd",
                cmd,
                "--cwd",
                cwd.to_string_lossy().as_ref(),
            ])
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .pre_exec(move || {
                if child_fd_copy != 3 {
                    let ret = libc::dup2(child_fd_copy, 3);
                    if ret < 0 {
                        return Err(std::io::Error::last_os_error());
                    }
                    libc::close(child_fd_copy);
                }
                Ok(())
            })
            .spawn()
            .expect("failed to spawn external shim")
    };

    unsafe {
        libc::close(child_fd);
    }

    let mut channel = Channel::new(parent_sock);
    channel
        .set_read_timeout(Some(Duration::from_millis(500)))
        .unwrap();
    (process, channel)
}

fn pid_exists(pid: u32) -> bool {
    let output = ProcessCommand::new("ps")
        .args(["-o", "stat=", "-p", &pid.to_string()])
        .output()
        .expect("failed to run ps");
    if !output.status.success() {
        return false;
    }

    let stat = String::from_utf8_lossy(&output.stdout);
    let trimmed = stat.trim();
    !trimmed.is_empty() && !trimmed.starts_with('Z')
}

/// Wait for a Ready event (with timeout). Handles read timeouts by retrying.
fn wait_for_ready(ch: &mut Channel) -> bool {
    let deadline = std::time::Instant::now() + Duration::from_secs(30);
    loop {
        if std::time::Instant::now() > deadline {
            return false;
        }
        match ch.recv::<Event>() {
            Ok(Some(Event::Ready)) => return true,
            Ok(Some(Event::StateChanged { .. })) => continue,
            Ok(Some(_)) => continue,
            Ok(None) => return false,
            Err(e) if is_timeout_error(&e) => continue,
            Err(_) => return false,
        }
    }
}

/// Drain events until we get a Completion or timeout.
fn wait_for_completion(ch: &mut Channel, timeout: Duration) -> Option<Event> {
    let deadline = std::time::Instant::now() + timeout;
    loop {
        if std::time::Instant::now() > deadline {
            return None;
        }
        match ch.recv::<Event>() {
            Ok(Some(evt @ Event::Completion { .. })) => return Some(evt),
            Ok(Some(Event::Died { .. })) => return None,
            Ok(Some(_)) => continue,
            Ok(None) => return None,
            Err(e) if is_timeout_error(&e) => continue,
            Err(_) => return None,
        }
    }
}

/// Drain events until we get a specific event type or timeout.
fn wait_for_event<F>(ch: &mut Channel, timeout: Duration, matcher: F) -> Option<Event>
where
    F: Fn(&Event) -> bool,
{
    let deadline = std::time::Instant::now() + timeout;
    loop {
        if std::time::Instant::now() > deadline {
            return None;
        }
        match ch.recv::<Event>() {
            Ok(Some(ref evt)) if matcher(evt) => return Some(evt.clone()),
            Ok(Some(_)) => continue,
            Ok(None) => return None,
            Err(e) if is_timeout_error(&e) => continue,
            Err(_) => return None,
        }
    }
}

fn wait_for_pid_file(path: &std::path::Path, timeout: Duration) -> Option<u32> {
    let deadline = std::time::Instant::now() + timeout;
    while std::time::Instant::now() <= deadline {
        if let Ok(contents) = fs::read_to_string(path) {
            if let Ok(pid) = contents.trim().parse::<u32>() {
                return Some(pid);
            }
        }
        std::thread::sleep(Duration::from_millis(100));
    }
    None
}

fn wait_for_pid_exit(pid: u32, timeout: Duration) -> bool {
    let deadline = std::time::Instant::now() + timeout;
    while std::time::Instant::now() <= deadline {
        if !pid_exists(pid) {
            return true;
        }
        std::thread::sleep(Duration::from_millis(100));
    }
    false
}

// -----------------------------------------------------------------------
// Tests — all gated behind shim-integration
// -----------------------------------------------------------------------

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_echo_hello() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo hello".into(),
        message_id: None,
    })
    .unwrap();

    let evt = wait_for_completion(&mut ch, Duration::from_secs(10));
    assert!(evt.is_some(), "did not receive Completion event");
    if let Some(Event::Completion { response, .. }) = evt {
        assert!(
            response.contains("hello"),
            "response should contain 'hello', got: {response}"
        );
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_exit_produces_died() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "exit".into(),
        message_id: None,
    })
    .unwrap();

    let evt = wait_for_event(&mut ch, Duration::from_secs(10), |e| {
        matches!(e, Event::Died { .. })
    });
    assert!(evt.is_some(), "did not receive Died event");
    // The PTY reader detects death via EOF, not waitpid, so exit_code
    // is None (the child process status isn't collected by the shim).
    if let Some(Event::Died { exit_code, .. }) = evt {
        assert!(
            exit_code.is_none() || exit_code == Some(0),
            "exit_code should be None or 0, got: {:?}",
            exit_code,
        );
    }
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_capture_screen() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    ch.send(&Command::CaptureScreen {
        last_n_lines: Some(10),
    })
    .unwrap();

    let evt = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(e, Event::ScreenCapture { .. })
    });
    assert!(evt.is_some(), "did not receive ScreenCapture event");
    if let Some(Event::ScreenCapture { content, .. }) = evt {
        assert!(!content.is_empty(), "screen capture should not be empty");
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_resize_no_error() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Resize should not produce an error event
    ch.send(&Command::Resize {
        rows: 40,
        cols: 120,
    })
    .unwrap();

    // Give it a moment, then verify with a Ping that the shim is still alive
    ch.send(&Command::Ping).unwrap();
    let evt = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(e, Event::Pong)
    });
    assert!(
        evt.is_some(),
        "shim should still be responsive after resize"
    );

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_ping_pong() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    ch.send(&Command::Ping).unwrap();
    let evt = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(e, Event::Pong)
    });
    assert!(evt.is_some(), "did not receive Pong event");

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_get_state() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    ch.send(&Command::GetState).unwrap();
    let evt = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(e, Event::State { .. })
    });
    assert!(evt.is_some(), "did not receive State event");
    if let Some(Event::State { state, .. }) = evt {
        assert_eq!(state, ShimState::Idle, "state should be Idle after ready");
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_send_while_working_returns_error() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Send a command that takes a moment
    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "sleep 2".into(),
        message_id: None,
    })
    .unwrap();

    // Wait for the Working state transition
    let working = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(
            e,
            Event::StateChanged {
                to: ShimState::Working,
                ..
            }
        )
    });
    assert!(working.is_some(), "should transition to Working");

    // Now try to send another message — should get Error
    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo should fail".into(),
        message_id: None,
    })
    .unwrap();

    let err = wait_for_event(&mut ch, Duration::from_secs(3), |e| {
        matches!(e, Event::Error { .. })
    });
    assert!(
        err.is_some(),
        "should receive Error when sending while working"
    );

    // Wait for the original command to complete
    let _ = wait_for_completion(&mut ch, Duration::from_secs(10));
    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
#[cfg_attr(
    feature = "shim-integration",
    ignore = "known regression: unexpected shim death may leave descendants alive"
)]
fn shim_death_reaps_background_descendants() {
    let tmp = tempfile::tempdir().unwrap();
    let pid_file = tmp.path().join("descendant.pid");
    let launcher = tmp.path().join("launch-descendant.py");

    fs::write(
        &launcher,
        format!(
            "import os\n\
import subprocess\n\
\n\
pid_file = r\"{}\"\n\
child_cmd = 'trap \"\" TERM HUP INT; echo $$ > \"{}\"; while :; do sleep 1; done'\n\
subprocess.Popen(['/bin/sh', '-c', child_cmd])\n\
os.execvp('bash', ['bash', '--noprofile', '--norc', '-i'])\n",
            pid_file.display(),
            pid_file.display()
        ),
    )
    .unwrap();

    let cmd = format!(
        "env BASH_ENV=/dev/null HOME='{}' python3 '{}'",
        tmp.path().display(),
        launcher.display()
    );

    let (mut shim, mut ch) = spawn_external_shim(&cmd, tmp.path());
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    let descendant_pid = wait_for_pid_file(&pid_file, Duration::from_secs(10))
        .expect("background descendant pid should be recorded");
    assert!(
        pid_exists(descendant_pid),
        "background descendant should be alive before shim death"
    );

    shim.kill().unwrap();
    let _ = shim.wait();

    assert!(
        wait_for_pid_exit(descendant_pid, Duration::from_secs(10)),
        "background descendant should be reaped after shim death"
    );
}

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_sleep_shows_working_then_idle() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "sleep 1 && echo done".into(),
        message_id: None,
    })
    .unwrap();

    // Should see Working state
    let working = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(
            e,
            Event::StateChanged {
                to: ShimState::Working,
                ..
            }
        )
    });
    assert!(working.is_some(), "should transition to Working");

    // Then completion (which implies transition back to Idle)
    let completion = wait_for_completion(&mut ch, Duration::from_secs(15));
    assert!(completion.is_some(), "should get completion after sleep");

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -----------------------------------------------------------------------
// E2E validation tests — Task #343
// -----------------------------------------------------------------------

/// Helper: spawn a shim with custom args in a background thread.
fn spawn_shim_with_args(args: ShimArgs) -> Channel {
    let (parent_sock, child_sock) = protocol::socketpair().unwrap();
    let channel = Channel::new(child_sock);

    std::thread::spawn(move || {
        crate::shim::runtime::run(args, channel).ok();
    });

    let mut ch = Channel::new(parent_sock);
    ch.set_read_timeout(Some(Duration::from_millis(500)))
        .unwrap();
    ch
}

/// Helper: spawn a shim with a PTY log path.
fn spawn_bash_shim_with_log(id: &str, log_path: PathBuf) -> Channel {
    let args = ShimArgs {
        id: id.into(),
        agent_type: crate::shim::classifier::AgentType::Generic,
        cmd: "bash --norc --noprofile".into(),
        cwd: PathBuf::from("/tmp"),
        rows: 24,
        cols: 80,
        pty_log_path: Some(log_path),
        graceful_shutdown_timeout_secs: 5,
        auto_commit_on_restart: true,
    };
    spawn_shim_with_args(args)
}

/// Helper: spawn a named shim (for multi-agent tests).
fn spawn_named_bash_shim(id: &str) -> Channel {
    let args = ShimArgs {
        id: id.into(),
        agent_type: crate::shim::classifier::AgentType::Generic,
        cmd: "bash --norc --noprofile".into(),
        cwd: PathBuf::from("/tmp"),
        rows: 24,
        cols: 80,
        pty_log_path: None,
        graceful_shutdown_timeout_secs: 5,
        auto_commit_on_restart: true,
    };
    spawn_shim_with_args(args)
}

/// Collect all events until timeout, returning them as a Vec.
fn collect_events(ch: &mut Channel, timeout: Duration) -> Vec<Event> {
    let deadline = std::time::Instant::now() + timeout;
    let mut events = Vec::new();
    loop {
        if std::time::Instant::now() > deadline {
            break;
        }
        match ch.recv::<Event>() {
            Ok(Some(evt)) => events.push(evt),
            Ok(None) => break,
            Err(e) if is_timeout_error(&e) => continue,
            Err(_) => break,
        }
    }
    events
}

// -- Test: Multiple concurrent agents --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_multiple_agents_concurrent() {
    // Spawn 3 independent shim agents
    let mut ch1 = spawn_named_bash_shim("agent-1");
    let mut ch2 = spawn_named_bash_shim("agent-2");
    let mut ch3 = spawn_named_bash_shim("agent-3");

    // All three should become ready independently
    assert!(wait_for_ready(&mut ch1), "agent-1 did not become ready");
    assert!(wait_for_ready(&mut ch2), "agent-2 did not become ready");
    assert!(wait_for_ready(&mut ch3), "agent-3 did not become ready");

    // Send different commands to each
    ch1.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo AGENT_ONE".into(),
        message_id: Some("msg-a1".into()),
    })
    .unwrap();

    ch2.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo AGENT_TWO".into(),
        message_id: Some("msg-a2".into()),
    })
    .unwrap();

    ch3.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo AGENT_THREE".into(),
        message_id: Some("msg-a3".into()),
    })
    .unwrap();

    // Each should produce its own completion
    let c1 = wait_for_completion(&mut ch1, Duration::from_secs(10));
    let c2 = wait_for_completion(&mut ch2, Duration::from_secs(10));
    let c3 = wait_for_completion(&mut ch3, Duration::from_secs(10));

    assert!(c1.is_some(), "agent-1 should complete");
    assert!(c2.is_some(), "agent-2 should complete");
    assert!(c3.is_some(), "agent-3 should complete");

    if let Some(Event::Completion { response, .. }) = c1 {
        assert!(
            response.contains("AGENT_ONE"),
            "agent-1 response should contain AGENT_ONE, got: {response}"
        );
    }
    if let Some(Event::Completion { response, .. }) = c2 {
        assert!(
            response.contains("AGENT_TWO"),
            "agent-2 response should contain AGENT_TWO, got: {response}"
        );
    }
    if let Some(Event::Completion { response, .. }) = c3 {
        assert!(
            response.contains("AGENT_THREE"),
            "agent-3 response should contain AGENT_THREE, got: {response}"
        );
    }

    // Verify independent state tracking via GetState
    for ch in [&mut ch1, &mut ch2, &mut ch3] {
        ch.send(&Command::GetState).unwrap();
        let state_evt = wait_for_event(ch, Duration::from_secs(5), |e| {
            matches!(e, Event::State { .. })
        });
        if let Some(Event::State { state, .. }) = state_evt {
            assert_eq!(state, ShimState::Idle, "all agents should be Idle");
        }
    }

    // Clean shutdown
    ch1.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
    ch2.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
    ch3.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -- Test: Crash recovery (kill child process) --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_crash_recovery_kill_child() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Kill the bash process from within (simulates crash)
    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "kill -9 $$".into(),
        message_id: None,
    })
    .unwrap();

    // Should receive a Died event
    let evt = wait_for_event(&mut ch, Duration::from_secs(10), |e| {
        matches!(e, Event::Died { .. })
    });
    assert!(evt.is_some(), "should receive Died event after kill -9");

    // Should also have transitioned to Dead state
    if let Some(Event::Died { last_lines, .. }) = evt {
        // last_lines should be non-empty (captures terminal state)
        assert!(
            !last_lines.trim().is_empty() || last_lines.is_empty(),
            "last_lines should be present (may be empty if process died instantly)"
        );
    }
}

// -- Test: Context exhaustion detection --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_context_exhaustion_detection() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Echo a known exhaustion pattern — the classifier checks screen content
    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo 'Error: context window is full'".into(),
        message_id: None,
    })
    .unwrap();

    // Should receive a ContextExhausted event (classifier detects the pattern)
    let evt = wait_for_event(&mut ch, Duration::from_secs(10), |e| {
        matches!(e, Event::ContextExhausted { .. })
    });
    assert!(
        evt.is_some(),
        "should receive ContextExhausted event when exhaustion text appears"
    );

    if let Some(Event::ContextExhausted { message, .. }) = evt {
        assert!(
            !message.is_empty(),
            "exhaustion message should not be empty"
        );
    }
}

// -- Test: Screen capture with cursor position --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_screen_capture_full() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Send full screen capture (no line limit)
    ch.send(&Command::CaptureScreen { last_n_lines: None })
        .unwrap();

    let evt = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(e, Event::ScreenCapture { .. })
    });
    assert!(evt.is_some(), "did not receive ScreenCapture event");
    if let Some(Event::ScreenCapture {
        content,
        cursor_row,
        cursor_col,
    }) = evt
    {
        assert!(
            !content.is_empty(),
            "full screen capture should not be empty"
        );
        // Cursor should be at a valid position
        assert!(cursor_row < 24, "cursor_row out of bounds: {cursor_row}");
        assert!(cursor_col < 80, "cursor_col out of bounds: {cursor_col}");
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -- Test: Graceful shutdown --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_graceful_shutdown() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Send a command so the agent is doing something
    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo 'before shutdown'".into(),
        message_id: None,
    })
    .unwrap();

    // Wait for completion
    let _ = wait_for_completion(&mut ch, Duration::from_secs(5));

    // Now request graceful shutdown
    ch.send(&Command::Shutdown {
        timeout_secs: 5,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();

    // After shutdown, the channel should close (recv returns None or Died)
    let events = collect_events(&mut ch, Duration::from_secs(10));
    // We may get StateChanged(→Dead), Died, or just EOF
    // The key assertion: the shim terminates without error
    let has_terminal = events.iter().any(|e| {
        matches!(
            e,
            Event::Died { .. }
                | Event::StateChanged {
                    to: ShimState::Dead,
                    ..
                }
        )
    }) || events.is_empty(); // EOF means clean exit
    assert!(
        has_terminal,
        "shutdown should result in clean termination, got: {:?}",
        events
    );
}

// -- Test: PTY log writing --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_pty_log_writing() {
    let tmp = tempfile::tempdir().unwrap();
    let log_path = tmp.path().join("test-agent.pty.log");

    let mut ch = spawn_bash_shim_with_log("log-agent", log_path.clone());
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Log file should exist after shim starts
    assert!(log_path.exists(), "PTY log file should exist after spawn");

    // Send a distinctive command
    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "echo PTY_LOG_MARKER_12345".into(),
        message_id: None,
    })
    .unwrap();

    let _ = wait_for_completion(&mut ch, Duration::from_secs(10));

    // Read the log file — it should contain raw PTY output
    let log_content = std::fs::read_to_string(&log_path).unwrap();
    assert!(
        !log_content.is_empty(),
        "PTY log should not be empty after commands"
    );
    assert!(
        log_content.contains("PTY_LOG_MARKER_12345"),
        "PTY log should contain command output, got: {}",
        &log_content[..log_content.len().min(200)]
    );

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -- Test: Message ID roundtrip --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_message_id_roundtrip() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    let msg_id = "test-msg-42";
    ch.send(&Command::SendMessage {
        from: "orchestrator".into(),
        body: "echo tracked".into(),
        message_id: Some(msg_id.into()),
    })
    .unwrap();

    let evt = wait_for_completion(&mut ch, Duration::from_secs(10));
    assert!(evt.is_some(), "should receive Completion");
    if let Some(Event::Completion { message_id, .. }) = evt {
        assert_eq!(
            message_id.as_deref(),
            Some(msg_id),
            "message_id should roundtrip through completion"
        );
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -- Test: Kill command --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_kill_command() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Kill should terminate immediately
    ch.send(&Command::Kill).unwrap();

    // Should get Died or EOF
    let events = collect_events(&mut ch, Duration::from_secs(10));
    let has_death = events.iter().any(|e| matches!(e, Event::Died { .. })) || events.is_empty(); // EOF = channel closed
    assert!(has_death, "Kill should terminate the shim");
}

// -- Test: State transition cycle (Starting → Idle → Working → Idle) --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_full_state_cycle() {
    let mut ch = spawn_bash_shim();

    // Phase 1: Starting → Idle (Ready event)
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Phase 2: Idle → Working
    ch.send(&Command::SendMessage {
        from: "user".into(),
        body: "sleep 1 && echo cycle_done".into(),
        message_id: None,
    })
    .unwrap();

    let working = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(
            e,
            Event::StateChanged {
                to: ShimState::Working,
                ..
            }
        )
    });
    assert!(working.is_some(), "should transition to Working");

    // Check from field
    if let Some(Event::StateChanged { from, to, .. }) = working {
        assert_eq!(from, ShimState::Idle, "should transition from Idle");
        assert_eq!(to, ShimState::Working, "should transition to Working");
    }

    // Phase 3: Working → Idle (via Completion)
    let idle = wait_for_event(&mut ch, Duration::from_secs(15), |e| {
        matches!(
            e,
            Event::StateChanged {
                to: ShimState::Idle,
                ..
            }
        )
    });
    assert!(idle.is_some(), "should transition back to Idle");

    if let Some(Event::StateChanged { from, to, .. }) = idle {
        assert_eq!(from, ShimState::Working);
        assert_eq!(to, ShimState::Idle);
    }

    // Verify final state
    ch.send(&Command::GetState).unwrap();
    let state_evt = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(e, Event::State { .. })
    });
    if let Some(Event::State { state, since_secs }) = state_evt {
        assert_eq!(state, ShimState::Idle);
        assert!(since_secs < 30, "state should be recent");
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -- Test: Multiple sequential commands --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_sequential_commands() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Send 3 sequential commands, each waiting for completion
    for i in 1..=3 {
        let cmd_body = format!("echo SEQ_{i}");
        ch.send(&Command::SendMessage {
            from: "user".into(),
            body: cmd_body,
            message_id: Some(format!("seq-{i}")),
        })
        .unwrap();

        let evt = wait_for_completion(&mut ch, Duration::from_secs(10));
        assert!(evt.is_some(), "should receive Completion for command {i}");
        if let Some(Event::Completion {
            response,
            message_id,
            ..
        }) = evt
        {
            assert!(
                response.contains(&format!("SEQ_{i}")),
                "response {i} should contain SEQ_{i}, got: {response}"
            );
            assert_eq!(
                message_id.as_deref(),
                Some(format!("seq-{i}").as_str()),
                "message_id should match for command {i}"
            );
        }
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -- Test: Channel-based lifecycle (simulates AgentHandle pattern) --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_channel_lifecycle_e2e() {
    // Create a socketpair and connect a shim to one end
    let (parent_sock, child_sock) = protocol::socketpair().unwrap();
    let child_channel = Channel::new(child_sock);

    let args = ShimArgs {
        id: "lifecycle-test".into(),
        agent_type: crate::shim::classifier::AgentType::Generic,
        cmd: "bash --norc --noprofile".into(),
        cwd: PathBuf::from("/tmp"),
        rows: 24,
        cols: 80,
        pty_log_path: None,
        graceful_shutdown_timeout_secs: 5,
        auto_commit_on_restart: true,
    };

    std::thread::spawn(move || {
        crate::shim::runtime::run(args, child_channel).ok();
    });

    let mut ch = Channel::new(parent_sock);
    ch.set_read_timeout(Some(Duration::from_millis(500)))
        .unwrap();

    // Track state manually (as AgentHandle would)
    let mut state = ShimState::Starting;

    // Phase 1: Wait for Ready, tracking state changes
    let deadline = std::time::Instant::now() + Duration::from_secs(30);
    let mut got_ready = false;
    while std::time::Instant::now() < deadline {
        match ch.recv::<Event>() {
            Ok(Some(Event::Ready)) => {
                state = ShimState::Idle;
                got_ready = true;
                break;
            }
            Ok(Some(Event::StateChanged { to, .. })) => {
                state = to;
            }
            Ok(Some(_)) => continue,
            Ok(None) => break,
            Err(e) if is_timeout_error(&e) => continue,
            Err(_) => break,
        }
    }
    assert!(got_ready, "should receive Ready");
    assert_eq!(state, ShimState::Idle);

    // Phase 2: Send message, track Working state
    ch.send(&Command::SendMessage {
        from: "test".into(),
        body: "echo LIFECYCLE_TEST".into(),
        message_id: Some("lc-1".into()),
    })
    .unwrap();

    // Wait for Working transition
    let working = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(
            e,
            Event::StateChanged {
                to: ShimState::Working,
                ..
            }
        )
    });
    assert!(working.is_some());
    state = ShimState::Working;

    // Wait for completion and return to Idle
    let deadline = std::time::Instant::now() + Duration::from_secs(10);
    let mut got_completion = false;
    while std::time::Instant::now() < deadline {
        match ch.recv::<Event>() {
            Ok(Some(Event::StateChanged { to, .. })) => {
                state = to;
            }
            Ok(Some(Event::Completion { message_id, .. })) => {
                assert_eq!(message_id.as_deref(), Some("lc-1"));
                got_completion = true;
                if state == ShimState::Idle {
                    break;
                }
            }
            Ok(Some(_)) => continue,
            Ok(None) => break,
            Err(e) if is_timeout_error(&e) => continue,
            Err(_) => break,
        }
    }
    assert!(got_completion, "should receive Completion");
    assert_eq!(state, ShimState::Idle);

    // Phase 3: Ping/Pong
    ch.send(&Command::Ping).unwrap();
    let pong = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
        matches!(e, Event::Pong)
    });
    assert!(pong.is_some(), "should receive Pong");

    // Phase 4: Graceful shutdown
    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}

// -- Test: Session state JSON roundtrip --

#[test]
fn shim_session_state_json_roundtrip() {
    // Test the shim state file format directly via JSON serialization
    // (the structs live in team::daemon::shim_state which is private,
    // so we validate the JSON format that save_shim_state produces)

    let state_json = serde_json::json!({
        "handles": {
            "eng-1": {
                "id": "eng-1",
                "agent_type": "claude",
                "agent_cmd": "claude --dangerously-skip-permissions",
                "work_dir": "/tmp/worktree/eng-1"
            },
            "eng-2": {
                "id": "eng-2",
                "agent_type": "codex",
                "agent_cmd": "codex",
                "work_dir": "/tmp/worktree/eng-2"
            },
            "eng-3": {
                "id": "eng-3",
                "agent_type": "generic",
                "agent_cmd": "bash",
                "work_dir": "/tmp/worktree/eng-3"
            }
        }
    });

    let tmp = tempfile::tempdir().unwrap();
    let state_dir = tmp.path().join(".batty");
    std::fs::create_dir_all(&state_dir).unwrap();
    let state_path = state_dir.join("shim_state.json");

    // Write JSON to disk
    let json_str = serde_json::to_string_pretty(&state_json).unwrap();
    std::fs::write(&state_path, &json_str).unwrap();

    // Read back and verify structure
    let content = std::fs::read_to_string(&state_path).unwrap();
    let loaded: serde_json::Value = serde_json::from_str(&content).unwrap();

    let handles = loaded["handles"].as_object().unwrap();
    assert_eq!(handles.len(), 3);

    // Verify each handle roundtripped
    assert_eq!(handles["eng-1"]["agent_type"], "claude");
    assert_eq!(handles["eng-2"]["agent_type"], "codex");
    assert_eq!(handles["eng-3"]["agent_type"], "generic");
    assert_eq!(handles["eng-1"]["work_dir"], "/tmp/worktree/eng-1");
    assert_eq!(
        handles["eng-1"]["agent_cmd"],
        "claude --dangerously-skip-permissions"
    );
}

// -- Test: Multiple pings --

#[test]
#[cfg_attr(not(feature = "shim-integration"), ignore)]
fn shim_multiple_pings() {
    let mut ch = spawn_bash_shim();
    assert!(wait_for_ready(&mut ch), "shim did not become ready");

    // Send 5 pings, each should produce a pong
    for i in 0..5 {
        ch.send(&Command::Ping).unwrap();
        let evt = wait_for_event(&mut ch, Duration::from_secs(5), |e| {
            matches!(e, Event::Pong)
        });
        assert!(evt.is_some(), "should receive Pong #{i}");
    }

    ch.send(&Command::Shutdown {
        timeout_secs: 2,
        reason: crate::shim::protocol::ShutdownReason::Requested,
    })
    .unwrap();
}