use std::collections::VecDeque;
use std::fs;
use std::io::{Read, Write as IoWrite};
use std::path::{Path, PathBuf};
use std::process::Command as ProcessCommand;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use std::time::Instant;
use anyhow::{Context, Result};
use portable_pty::{Child, CommandBuilder, PtySize};
use super::classifier::{self, AgentType, ScreenVerdict};
use super::common::{self, QueuedMessage};
use super::protocol::{Channel, Command, Event, ShimState};
use super::pty_log::PtyLogWriter;
use crate::prompt::strip_ansi;
const DEFAULT_ROWS: u16 = 50;
const DEFAULT_COLS: u16 = 220;
const SCROLLBACK_LINES: usize = 5000;
const POLL_INTERVAL_MS: u64 = 250;
const WORKING_DWELL_MS: u64 = 300;
const KIRO_IDLE_SETTLE_MS: u64 = 1200;
const READY_TIMEOUT_SECS: u64 = 120;
use common::MAX_QUEUE_DEPTH;
use common::SESSION_STATS_INTERVAL_SECS;
const PROCESS_EXIT_POLL_MS: u64 = 100;
const PARENT_DEATH_POLL_SECS: u64 = 1;
const GROUP_TERM_GRACE_SECS: u64 = 2;
pub(crate) const HANDOFF_FILE_NAME: &str = ".batty-handoff.md";
const AUTO_COMMIT_MESSAGE: &str = "wip: auto-save before restart [batty]";
const AUTO_COMMIT_TIMEOUT_SECS: u64 = 5;
pub(crate) fn preserve_handoff(
worktree: &Path,
task: &crate::task::Task,
recent_output: Option<&str>,
) -> Result<()> {
let branch_name = git_capture(worktree, &["branch", "--show-current"]).unwrap_or_default();
let last_commit = git_capture(worktree, &["rev-parse", "HEAD"]).unwrap_or_default();
let changed_files = summarize_changed_files(worktree);
let recent_commits = git_capture(worktree, &["log", "--oneline", "-5"]).unwrap_or_default();
let tests_run = recent_output
.map(extract_test_commands)
.unwrap_or_default()
.join("\n");
let progress_summary = summarize_recent_progress(worktree, task, recent_output);
let recent_activity = recent_output
.map(summarize_recent_activity)
.unwrap_or_default();
let handoff = format!(
"# Carry-Forward Summary\n## Task Spec\nTask #{}: {}\n\n{}\n\n## Work Completed So Far\n### Branch\n{}\n\n### Last Commit\n{}\n\n### Changed Files\n{}\n\n### Tests Run\n{}\n\n### Progress Summary\n{}\n\n### Recent Activity\n{}\n\n### Recent Commits\n{}\n\n## What Remains\n{}\n",
task.id,
task.title,
empty_section_fallback(&task.description),
empty_section_fallback(&branch_name),
empty_section_fallback(&last_commit),
empty_section_fallback(&changed_files),
empty_section_fallback(&tests_run),
empty_section_fallback(&progress_summary),
empty_section_fallback(&recent_activity),
empty_section_fallback(&recent_commits),
handoff_remaining_work(task)
);
fs::write(worktree.join(HANDOFF_FILE_NAME), handoff)
.with_context(|| format!("failed to write handoff file in {}", worktree.display()))?;
Ok(())
}
fn git_capture(worktree: &Path, args: &[&str]) -> Result<String> {
let output = ProcessCommand::new("git")
.args(args)
.current_dir(worktree)
.env_remove("GIT_DIR")
.env_remove("GIT_WORK_TREE")
.output()
.with_context(|| {
format!(
"failed to run `git {}` in {}",
args.join(" "),
worktree.display()
)
})?;
if !output.status.success() {
anyhow::bail!(
"`git {}` failed in {}: {}",
args.join(" "),
worktree.display(),
String::from_utf8_lossy(&output.stderr).trim()
);
}
Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
}
fn empty_section_fallback(content: &str) -> &str {
if content.trim().is_empty() {
"(none)"
} else {
content
}
}
fn summarize_recent_activity(output: &str) -> String {
let cleaned = strip_ansi(output);
let lines: Vec<&str> = cleaned
.lines()
.map(str::trim_end)
.filter(|line| !line.trim().is_empty())
.collect();
let start = lines.len().saturating_sub(40);
lines[start..].join("\n")
}
fn summarize_recent_progress(
worktree: &Path,
task: &crate::task::Task,
recent_output: Option<&str>,
) -> String {
let mut summary = summarize_progress_from_event_log(worktree, task);
if summary.trim().is_empty() {
summary = recent_output
.map(summarize_recent_activity)
.unwrap_or_default();
}
truncate_to_word_limit(&summary, 500)
}
fn summarize_progress_from_event_log(worktree: &Path, task: &crate::task::Task) -> String {
let events_path = crate::team::team_events_path(worktree);
let Ok(events) = crate::team::events::read_events(&events_path) else {
return String::new();
};
let task_id = task.id.to_string();
let claimed_by = task.claimed_by.as_deref();
let mut lines = Vec::new();
for event in events.iter().rev() {
let matches_task = event.task.as_deref() == Some(task_id.as_str());
let matches_role = claimed_by.is_some() && event.role.as_deref() == claimed_by;
if !matches_task && !matches_role {
continue;
}
let mut line = event.event.clone();
if let Some(reason) = event
.reason
.as_deref()
.filter(|reason| !reason.trim().is_empty())
{
line.push_str(": ");
line.push_str(reason.trim());
} else if let Some(details) = event
.details
.as_deref()
.filter(|details| !details.trim().is_empty())
{
line.push_str(": ");
line.push_str(details.trim());
} else if let Some(error) = event
.error
.as_deref()
.filter(|error| !error.trim().is_empty())
{
line.push_str(": ");
line.push_str(error.trim());
}
lines.push(line);
if lines.len() >= 12 {
break;
}
}
lines.reverse();
lines.join("\n")
}
fn truncate_to_word_limit(input: &str, word_limit: usize) -> String {
let words: Vec<&str> = input.split_whitespace().collect();
if words.len() <= word_limit {
return input.trim().to_string();
}
let truncated = words[..word_limit].join(" ");
format!("{truncated} ...")
}
fn summarize_changed_files(worktree: &Path) -> String {
let mut files = Vec::new();
for args in [
&["diff", "--name-only"] as &[&str],
&["diff", "--cached", "--name-only"],
&["ls-files", "--others", "--exclude-standard"],
] {
let Ok(output) = git_capture(worktree, args) else {
continue;
};
for line in output.lines() {
let trimmed = line.trim();
if !trimmed.is_empty() && !files.iter().any(|existing| existing == trimmed) {
files.push(trimmed.to_string());
}
}
}
files.join("\n")
}
fn handoff_remaining_work(task: &crate::task::Task) -> &str {
task.next_action
.as_deref()
.filter(|next| !next.trim().is_empty())
.unwrap_or("Continue from the current worktree state, verify acceptance criteria, and finish the task without redoing completed work.")
}
fn extract_test_commands(output: &str) -> Vec<String> {
let cleaned = strip_ansi(output);
let mut commands = Vec::new();
for line in cleaned.lines() {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let lower = trimmed.to_ascii_lowercase();
if (lower.contains("cargo test")
|| lower.contains("cargo nextest")
|| lower.contains("pytest")
|| lower.contains("npm test")
|| lower.contains("pnpm test")
|| lower.contains("yarn test")
|| lower.contains("go test")
|| lower.contains("bundle exec rspec")
|| lower.contains("mix test"))
&& !commands.iter().any(|existing| existing == trimmed)
{
commands.push(trimmed.to_string());
}
}
commands
}
fn format_injected_message(sender: &str, body: &str) -> String {
common::format_injected_message(sender, body)
}
fn shell_single_quote(input: &str) -> String {
input.replace('\'', "'\\''")
}
fn build_supervised_agent_command(command: &str, shim_pid: u32) -> String {
let escaped_command = shell_single_quote(command);
format!(
"shim_pid={shim_pid}; \
agent_root_pid=$$; \
agent_pgid=$$; \
setsid sh -c ' \
shim_pid=\"$1\"; \
agent_pgid=\"$2\"; \
agent_root_pid=\"$3\"; \
collect_descendants() {{ \
parent_pid=\"$1\"; \
for child_pid in $(pgrep -P \"$parent_pid\" 2>/dev/null); do \
printf \"%s\\n\" \"$child_pid\"; \
collect_descendants \"$child_pid\"; \
done; \
}}; \
while kill -0 \"$shim_pid\" 2>/dev/null; do sleep {PARENT_DEATH_POLL_SECS}; done; \
descendant_pids=$(collect_descendants \"$agent_root_pid\"); \
kill -TERM -- -\"$agent_pgid\" >/dev/null 2>&1 || true; \
for descendant_pid in $descendant_pids; do kill -TERM \"$descendant_pid\" >/dev/null 2>&1 || true; done; \
sleep {GROUP_TERM_GRACE_SECS}; \
kill -KILL -- -\"$agent_pgid\" >/dev/null 2>&1 || true; \
for descendant_pid in $descendant_pids; do kill -KILL \"$descendant_pid\" >/dev/null 2>&1 || true; done \
' _ \"$shim_pid\" \"$agent_pgid\" \"$agent_root_pid\" >/dev/null 2>&1 < /dev/null & \
exec bash -lc '{escaped_command}'"
)
}
#[cfg(unix)]
fn signal_process_group(child: &dyn Child, signal: libc::c_int) -> std::io::Result<()> {
let pid = child
.process_id()
.ok_or_else(|| std::io::Error::other("child process id unavailable"))?;
let result = unsafe { libc::killpg(pid as libc::pid_t, signal) };
if result == 0 {
Ok(())
} else {
Err(std::io::Error::last_os_error())
}
}
fn terminate_agent_group(
child: &mut Box<dyn Child + Send + Sync>,
sigterm_grace: Duration,
) -> std::io::Result<()> {
#[cfg(unix)]
{
signal_process_group(child.as_ref(), libc::SIGTERM)?;
let deadline = Instant::now() + sigterm_grace;
while Instant::now() <= deadline {
if let Ok(Some(_)) = child.try_wait() {
return Ok(());
}
thread::sleep(Duration::from_millis(PROCESS_EXIT_POLL_MS));
}
signal_process_group(child.as_ref(), libc::SIGKILL)?;
return Ok(());
}
#[allow(unreachable_code)]
child.kill()
}
fn graceful_shutdown_timeout() -> Duration {
let secs = std::env::var("BATTY_GRACEFUL_SHUTDOWN_TIMEOUT_SECS")
.ok()
.and_then(|value| value.parse::<u64>().ok())
.unwrap_or(AUTO_COMMIT_TIMEOUT_SECS);
Duration::from_secs(secs)
}
fn auto_commit_on_restart_enabled() -> bool {
std::env::var("BATTY_AUTO_COMMIT_ON_RESTART")
.map(|value| !matches!(value.as_str(), "0" | "false" | "FALSE"))
.unwrap_or(true)
}
fn preserve_work_before_kill_with<F>(
worktree_path: &Path,
timeout: Duration,
enabled: bool,
commit_fn: F,
) -> Result<bool>
where
F: FnOnce(PathBuf) -> Result<bool> + Send + 'static,
{
if !enabled {
return Ok(false);
}
let (tx, rx) = mpsc::channel();
let path = worktree_path.to_path_buf();
thread::spawn(move || {
let _ = tx.send(commit_fn(path));
});
match rx.recv_timeout(timeout) {
Ok(result) => result,
Err(mpsc::RecvTimeoutError::Timeout) => Ok(false),
Err(mpsc::RecvTimeoutError::Disconnected) => Ok(false),
}
}
pub(crate) fn preserve_work_before_kill(worktree_path: &Path) -> Result<bool> {
let timeout = graceful_shutdown_timeout();
preserve_work_before_kill_with(
worktree_path,
timeout,
auto_commit_on_restart_enabled(),
move |path| {
crate::team::git_cmd::auto_commit_if_dirty(&path, AUTO_COMMIT_MESSAGE, timeout)
.map_err(anyhow::Error::from)
},
)
}
fn pty_write_paced(
pty_writer: &Arc<Mutex<Box<dyn std::io::Write + Send>>>,
agent_type: AgentType,
body: &[u8],
enter: &[u8],
) -> std::io::Result<()> {
match agent_type {
AgentType::Generic => {
let mut writer = pty_writer.lock().unwrap();
writer.write_all(body)?;
writer.write_all(enter)?;
writer.flush()?;
}
_ => {
let mut writer = pty_writer.lock().unwrap();
writer.write_all(b"\x1b[200~")?;
writer.write_all(body)?;
writer.write_all(b"\x1b[201~")?;
writer.flush()?;
drop(writer);
std::thread::sleep(std::time::Duration::from_millis(200));
let mut writer = pty_writer.lock().unwrap();
writer.write_all(enter)?;
writer.flush()?;
}
}
Ok(())
}
fn enter_seq(agent_type: AgentType) -> &'static str {
match agent_type {
AgentType::Generic => "\n",
_ => "\r", }
}
#[derive(Debug, Clone)]
pub struct ShimArgs {
pub id: String,
pub agent_type: AgentType,
pub cmd: String,
pub cwd: PathBuf,
pub rows: u16,
pub cols: u16,
pub pty_log_path: Option<PathBuf>,
pub graceful_shutdown_timeout_secs: u64,
pub auto_commit_on_restart: bool,
}
impl ShimArgs {
fn preserve_work_before_kill(&self, worktree_path: &Path) -> Result<bool> {
if !self.auto_commit_on_restart {
return Ok(false);
}
let status = ProcessCommand::new("git")
.arg("-C")
.arg(worktree_path)
.args(["status", "--porcelain"])
.output()
.with_context(|| {
format!(
"failed to inspect git status in {}",
worktree_path.display()
)
})?;
if !status.status.success() {
anyhow::bail!("git status failed in {}", worktree_path.display());
}
let dirty = String::from_utf8_lossy(&status.stdout)
.lines()
.any(|line| !line.starts_with("?? .batty/"));
if !dirty {
return Ok(false);
}
let timeout = Duration::from_secs(self.graceful_shutdown_timeout_secs);
run_git_preserve_with_timeout(worktree_path, &["add", "-A"], timeout)?;
run_git_preserve_with_timeout(
worktree_path,
&["commit", "-m", "wip: auto-save before restart [batty]"],
timeout,
)?;
Ok(true)
}
}
fn run_git_preserve_with_timeout(
worktree_path: &Path,
args: &[&str],
timeout: Duration,
) -> Result<()> {
let mut child = ProcessCommand::new("git")
.arg("-C")
.arg(worktree_path)
.args(args)
.spawn()
.with_context(|| {
format!(
"failed to launch `git {}` in {}",
args.join(" "),
worktree_path.display()
)
})?;
let deadline = Instant::now() + timeout;
loop {
if let Some(status) = child.try_wait()? {
if status.success() {
return Ok(());
}
anyhow::bail!(
"`git {}` failed in {} with status {}",
args.join(" "),
worktree_path.display(),
status
);
}
if Instant::now() >= deadline {
let _ = child.kill();
let _ = child.wait();
anyhow::bail!(
"`git {}` timed out after {}s in {}",
args.join(" "),
timeout.as_secs(),
worktree_path.display()
);
}
thread::sleep(Duration::from_millis(50));
}
}
struct ShimInner {
parser: vt100::Parser,
state: ShimState,
state_changed_at: Instant,
last_screen_hash: u64,
last_pty_output_at: Instant,
started_at: Instant,
active_session_started_at: Instant,
cumulative_output_bytes: u64,
pre_injection_content: String,
pending_message_id: Option<String>,
agent_type: AgentType,
message_queue: VecDeque<QueuedMessage>,
dialogs_dismissed: u8,
last_working_screen: String,
test_failure_iterations: u8,
context_approaching_emitted: bool,
}
impl ShimInner {
fn screen_contents(&self) -> String {
self.parser.screen().contents()
}
fn last_n_lines(&self, n: usize) -> String {
let content = self.parser.screen().contents();
let lines: Vec<&str> = content.lines().collect();
let start = lines.len().saturating_sub(n);
lines[start..].join("\n")
}
fn cursor_position(&self) -> (u16, u16) {
self.parser.screen().cursor_position()
}
}
fn content_hash(s: &str) -> u64 {
let mut hash: u64 = 0xcbf29ce484222325;
for byte in s.bytes() {
hash ^= byte as u64;
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}
pub fn run(args: ShimArgs, channel: Channel) -> Result<()> {
let rows = if args.rows > 0 {
args.rows
} else {
DEFAULT_ROWS
};
let cols = if args.cols > 0 {
args.cols
} else {
DEFAULT_COLS
};
let pty_system = portable_pty::native_pty_system();
let pty_pair = pty_system
.openpty(PtySize {
rows,
cols,
pixel_width: 0,
pixel_height: 0,
})
.context("failed to create PTY")?;
let shim_pid = std::process::id();
let supervised_cmd = build_supervised_agent_command(&args.cmd, shim_pid);
let mut cmd = CommandBuilder::new("bash");
cmd.args(["-lc", &supervised_cmd]);
cmd.cwd(&args.cwd);
cmd.env_remove("CLAUDECODE"); cmd.env("TERM", "xterm-256color");
cmd.env("COLORTERM", "truecolor");
let mut child = pty_pair
.slave
.spawn_command(cmd)
.context("failed to spawn agent CLI")?;
drop(pty_pair.slave);
let mut pty_reader = pty_pair
.master
.try_clone_reader()
.context("failed to clone PTY reader")?;
let pty_writer = pty_pair
.master
.take_writer()
.context("failed to take PTY writer")?;
let inner = Arc::new(Mutex::new(ShimInner {
parser: vt100::Parser::new(rows, cols, SCROLLBACK_LINES),
state: ShimState::Starting,
state_changed_at: Instant::now(),
last_screen_hash: 0,
last_pty_output_at: Instant::now(),
started_at: Instant::now(),
active_session_started_at: Instant::now(),
cumulative_output_bytes: 0,
pre_injection_content: String::new(),
pending_message_id: None,
agent_type: args.agent_type,
message_queue: VecDeque::new(),
dialogs_dismissed: 0,
last_working_screen: String::new(),
test_failure_iterations: 0,
context_approaching_emitted: false,
}));
let pty_log: Option<Mutex<PtyLogWriter>> = args
.pty_log_path
.as_deref()
.map(|p| PtyLogWriter::new(p).context("failed to create PTY log"))
.transpose()?
.map(Mutex::new);
let pty_log = pty_log.map(Arc::new);
let pty_writer = Arc::new(Mutex::new(pty_writer));
let mut cmd_channel = channel;
let mut evt_channel = cmd_channel.try_clone().context("failed to clone channel")?;
let inner_pty = Arc::clone(&inner);
let log_handle = pty_log.clone();
let pty_writer_pty = Arc::clone(&pty_writer);
let pty_handle = std::thread::spawn(move || {
let mut buf = [0u8; 4096];
loop {
match pty_reader.read(&mut buf) {
Ok(0) => break, Ok(n) => {
if let Some(ref log) = log_handle {
let _ = log.lock().unwrap().write(&buf[..n]);
}
let mut inner = inner_pty.lock().unwrap();
inner.last_pty_output_at = Instant::now();
inner.cumulative_output_bytes =
inner.cumulative_output_bytes.saturating_add(n as u64);
inner.parser.process(&buf[..n]);
let content = inner.parser.screen().contents();
let hash = content_hash(&content);
if hash == inner.last_screen_hash {
continue; }
inner.last_screen_hash = hash;
let verdict = classifier::classify(inner.agent_type, inner.parser.screen());
let old_state = inner.state;
if !inner.context_approaching_emitted
&& common::detect_context_approaching_limit(&content)
{
inner.context_approaching_emitted = true;
drop(inner);
let _ = evt_channel.send(&Event::ContextApproaching {
message: "Screen output contains context-pressure signals".into(),
input_tokens: 0,
output_tokens: 0,
});
continue;
}
if old_state == ShimState::Working {
inner.last_working_screen = content.clone();
}
let working_too_short = old_state == ShimState::Working
&& inner.state_changed_at.elapsed().as_millis() < WORKING_DWELL_MS as u128;
let new_state = match (old_state, verdict) {
(ShimState::Starting, ScreenVerdict::AgentIdle) => Some(ShimState::Idle),
(ShimState::Idle, ScreenVerdict::AgentIdle) => None,
(ShimState::Working, ScreenVerdict::AgentIdle) if working_too_short => None,
(ShimState::Working, ScreenVerdict::AgentIdle)
if inner.agent_type == AgentType::Kiro =>
{
None
}
(ShimState::Working, ScreenVerdict::AgentIdle) => Some(ShimState::Idle),
(ShimState::Working, ScreenVerdict::AgentWorking) => None,
(_, ScreenVerdict::ContextExhausted) => Some(ShimState::ContextExhausted),
(_, ScreenVerdict::Unknown) => None,
(ShimState::Idle, ScreenVerdict::AgentWorking) => Some(ShimState::Working),
(ShimState::Starting, ScreenVerdict::AgentWorking) => {
Some(ShimState::Working)
}
_ => None,
};
if let Some(new) = new_state {
let summary = inner.last_n_lines(5);
inner.state = new;
inner.state_changed_at = Instant::now();
let pre_content = inner.pre_injection_content.clone();
let current_content = inner.screen_contents();
let working_screen = inner.last_working_screen.clone();
let msg_id = inner.pending_message_id.take();
let response =
completion_response(&pre_content, ¤t_content, &working_screen);
if old_state == ShimState::Working && new == ShimState::Idle {
if let Some(followup) = common::detect_test_failure_followup(
&response,
inner.test_failure_iterations,
) {
inner.test_failure_iterations = followup.next_iteration_count;
inner.pre_injection_content = inner.screen_contents();
inner.pending_message_id = None;
inner.state = ShimState::Working;
inner.state_changed_at = Instant::now();
inner.active_session_started_at = Instant::now();
let agent_type_for_enter = inner.agent_type;
drop(inner);
let enter = enter_seq(agent_type_for_enter);
if let Err(e) = pty_write_paced(
&pty_writer_pty,
agent_type_for_enter,
format_injected_message("batty", &followup.body).as_bytes(),
enter.as_bytes(),
) {
let _ = evt_channel.send(&Event::Error {
command: "SendMessage".into(),
reason: format!(
"PTY write failed for test failure follow-up: {e}"
),
});
} else {
let _ = evt_channel.send(&Event::Warning {
message: followup.notice,
idle_secs: None,
});
}
continue;
}
inner.test_failure_iterations = 0;
}
let drain_errors =
if new == ShimState::Dead || new == ShimState::ContextExhausted {
drain_queue_errors(&mut inner.message_queue, new)
} else {
Vec::new()
};
let queued_msg = if old_state == ShimState::Working
&& new == ShimState::Idle
&& !inner.message_queue.is_empty()
{
inner.message_queue.pop_front()
} else {
None
};
if let Some(ref msg) = queued_msg {
inner.pre_injection_content = inner.screen_contents();
inner.pending_message_id = msg.message_id.clone();
inner.state = ShimState::Working;
inner.state_changed_at = Instant::now();
inner.active_session_started_at = Instant::now();
inner.test_failure_iterations = 0;
}
let queue_depth = inner.message_queue.len();
let agent_type_for_enter = inner.agent_type;
let queued_injected = queued_msg
.as_ref()
.map(|msg| format_injected_message(&msg.from, &msg.body));
drop(inner);
let events = build_transition_events(
old_state,
new,
&summary,
&pre_content,
¤t_content,
&working_screen,
msg_id,
);
for event in events {
if evt_channel.send(&event).is_err() {
return; }
}
for event in drain_errors {
if evt_channel.send(&event).is_err() {
return;
}
}
if let Some(msg) = queued_msg {
let enter = enter_seq(agent_type_for_enter);
let injected = queued_injected.as_deref().unwrap_or(msg.body.as_str());
if let Err(e) = pty_write_paced(
&pty_writer_pty,
agent_type_for_enter,
injected.as_bytes(),
enter.as_bytes(),
) {
let _ = evt_channel.send(&Event::Error {
command: "SendMessage".into(),
reason: format!("PTY write failed for queued message: {e}"),
});
}
let _ = evt_channel.send(&Event::StateChanged {
from: ShimState::Idle,
to: ShimState::Working,
summary: format!(
"delivering queued message ({} remaining)",
queue_depth
),
});
}
}
}
Err(_) => break, }
}
let mut inner = inner_pty.lock().unwrap();
let last_lines = inner.last_n_lines(10);
let old = inner.state;
inner.state = ShimState::Dead;
let drain_errors = drain_queue_errors(&mut inner.message_queue, ShimState::Dead);
drop(inner);
let _ = evt_channel.send(&Event::StateChanged {
from: old,
to: ShimState::Dead,
summary: last_lines.clone(),
});
let _ = evt_channel.send(&Event::Died {
exit_code: None,
last_lines,
});
for event in drain_errors {
let _ = evt_channel.send(&event);
}
});
let inner_idle = Arc::clone(&inner);
let pty_writer_idle = Arc::clone(&pty_writer);
let mut idle_channel = cmd_channel.try_clone().context("failed to clone channel")?;
std::thread::spawn(move || {
loop {
std::thread::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS));
let mut inner = inner_idle.lock().unwrap();
if inner.agent_type != AgentType::Kiro || inner.state != ShimState::Working {
continue;
}
if inner.last_pty_output_at.elapsed().as_millis() < KIRO_IDLE_SETTLE_MS as u128 {
continue;
}
if classifier::classify(inner.agent_type, inner.parser.screen())
!= ScreenVerdict::AgentIdle
{
continue;
}
let summary = inner.last_n_lines(5);
let pre_content = inner.pre_injection_content.clone();
let current_content = inner.screen_contents();
let working_screen = inner.last_working_screen.clone();
let msg_id = inner.pending_message_id.take();
let response = completion_response(&pre_content, ¤t_content, &working_screen);
if let Some(followup) =
common::detect_test_failure_followup(&response, inner.test_failure_iterations)
{
inner.test_failure_iterations = followup.next_iteration_count;
inner.pre_injection_content = inner.screen_contents();
inner.pending_message_id = None;
inner.state = ShimState::Working;
inner.state_changed_at = Instant::now();
inner.active_session_started_at = Instant::now();
let agent_type_for_enter = inner.agent_type;
drop(inner);
let enter = enter_seq(agent_type_for_enter);
if let Err(e) = pty_write_paced(
&pty_writer_idle,
agent_type_for_enter,
format_injected_message("batty", &followup.body).as_bytes(),
enter.as_bytes(),
) {
let _ = idle_channel.send(&Event::Error {
command: "SendMessage".into(),
reason: format!("PTY write failed for test failure follow-up: {e}"),
});
} else {
let _ = idle_channel.send(&Event::Warning {
message: followup.notice,
idle_secs: None,
});
}
continue;
}
inner.test_failure_iterations = 0;
inner.state = ShimState::Idle;
inner.state_changed_at = Instant::now();
let queued_msg = if !inner.message_queue.is_empty() {
inner.message_queue.pop_front()
} else {
None
};
if let Some(ref msg) = queued_msg {
inner.pre_injection_content = inner.screen_contents();
inner.pending_message_id = msg.message_id.clone();
inner.state = ShimState::Working;
inner.state_changed_at = Instant::now();
inner.active_session_started_at = Instant::now();
inner.test_failure_iterations = 0;
}
let queue_depth = inner.message_queue.len();
let agent_type_for_enter = inner.agent_type;
let queued_injected = queued_msg
.as_ref()
.map(|msg| format_injected_message(&msg.from, &msg.body));
drop(inner);
for event in build_transition_events(
ShimState::Working,
ShimState::Idle,
&summary,
&pre_content,
¤t_content,
&working_screen,
msg_id,
) {
if idle_channel.send(&event).is_err() {
return;
}
}
if let Some(msg) = queued_msg {
let enter = enter_seq(agent_type_for_enter);
let injected = queued_injected.as_deref().unwrap_or(msg.body.as_str());
if let Err(e) = pty_write_paced(
&pty_writer_idle,
agent_type_for_enter,
injected.as_bytes(),
enter.as_bytes(),
) {
let _ = idle_channel.send(&Event::Error {
command: "SendMessage".into(),
reason: format!("PTY write failed for queued message: {e}"),
});
continue;
}
let _ = idle_channel.send(&Event::StateChanged {
from: ShimState::Idle,
to: ShimState::Working,
summary: format!("delivering queued message ({} remaining)", queue_depth),
});
}
}
});
let inner_poll = Arc::clone(&inner);
let mut poll_channel = cmd_channel
.try_clone()
.context("failed to clone channel for poll thread")?;
std::thread::spawn(move || {
loop {
std::thread::sleep(std::time::Duration::from_secs(5));
let mut inner = inner_poll.lock().unwrap();
if inner.state != ShimState::Working {
continue;
}
if inner.last_pty_output_at.elapsed().as_secs() < 2 {
continue;
}
let verdict = classifier::classify(inner.agent_type, inner.parser.screen());
if verdict == classifier::ScreenVerdict::AgentIdle {
let summary = inner.last_n_lines(5);
inner.state = ShimState::Idle;
inner.state_changed_at = Instant::now();
drop(inner);
let _ = poll_channel.send(&Event::StateChanged {
from: ShimState::Working,
to: ShimState::Idle,
summary,
});
}
}
});
let inner_stats = Arc::clone(&inner);
let mut stats_channel = cmd_channel
.try_clone()
.context("failed to clone channel for stats thread")?;
std::thread::spawn(move || {
loop {
std::thread::sleep(Duration::from_secs(SESSION_STATS_INTERVAL_SECS));
let inner = inner_stats.lock().unwrap();
if inner.state == ShimState::Dead {
return;
}
let output_bytes = inner.cumulative_output_bytes;
let uptime_secs = match inner.state {
ShimState::Working | ShimState::ContextExhausted => {
inner.active_session_started_at.elapsed().as_secs()
}
_ => inner.started_at.elapsed().as_secs(),
};
drop(inner);
if stats_channel
.send(&Event::SessionStats {
output_bytes,
uptime_secs,
input_tokens: 0,
output_tokens: 0,
context_usage_pct: None,
})
.is_err()
{
return;
}
}
});
let inner_cmd = Arc::clone(&inner);
let start = Instant::now();
loop {
let mut inner = inner_cmd.lock().unwrap();
let state = inner.state;
match state {
ShimState::Starting => {
if inner.dialogs_dismissed < 10 {
let content = inner.screen_contents();
if classifier::detect_startup_dialog(&content) {
let attempt = inner.dialogs_dismissed + 1;
let enter = enter_seq(inner.agent_type);
inner.dialogs_dismissed = attempt;
drop(inner);
eprintln!(
"[shim {}] auto-dismissing startup dialog (attempt {attempt})",
args.id
);
let mut writer = pty_writer.lock().unwrap();
writer.write_all(enter.as_bytes()).ok();
writer.flush().ok();
std::thread::sleep(std::time::Duration::from_millis(POLL_INTERVAL_MS));
continue;
}
}
drop(inner);
if start.elapsed().as_secs() > READY_TIMEOUT_SECS {
let last = inner_cmd.lock().unwrap().last_n_lines(10);
cmd_channel.send(&Event::Error {
command: "startup".into(),
reason: format!(
"agent did not show prompt within {}s. Last lines:\n{}",
READY_TIMEOUT_SECS, last,
),
})?;
terminate_agent_group(&mut child, Duration::from_secs(GROUP_TERM_GRACE_SECS))
.ok();
return Ok(());
}
thread::sleep(Duration::from_millis(POLL_INTERVAL_MS));
}
ShimState::Dead => {
drop(inner);
return Ok(());
}
ShimState::Idle => {
drop(inner);
cmd_channel.send(&Event::Ready)?;
break;
}
_ => {
drop(inner);
if start.elapsed().as_secs() > READY_TIMEOUT_SECS {
let last = inner_cmd.lock().unwrap().last_n_lines(10);
cmd_channel.send(&Event::Error {
command: "startup".into(),
reason: format!(
"agent did not reach idle within {}s (state: {}). Last lines:\n{}",
READY_TIMEOUT_SECS, state, last,
),
})?;
terminate_agent_group(&mut child, Duration::from_secs(GROUP_TERM_GRACE_SECS))
.ok();
return Ok(());
}
thread::sleep(Duration::from_millis(POLL_INTERVAL_MS));
}
}
}
loop {
let cmd = match cmd_channel.recv::<Command>() {
Ok(Some(c)) => c,
Ok(None) => {
eprintln!(
"[shim {}] orchestrator disconnected, shutting down",
args.id
);
terminate_agent_group(&mut child, Duration::from_secs(GROUP_TERM_GRACE_SECS)).ok();
break;
}
Err(e) => {
eprintln!("[shim {}] channel error: {e}", args.id);
terminate_agent_group(&mut child, Duration::from_secs(GROUP_TERM_GRACE_SECS)).ok();
break;
}
};
match cmd {
Command::SendMessage {
from,
body,
message_id,
} => {
let delivery_id = message_id.clone();
let mut inner = inner_cmd.lock().unwrap();
match inner.state {
ShimState::Idle => {
inner.pre_injection_content = inner.screen_contents();
inner.pending_message_id = message_id;
inner.test_failure_iterations = 0;
let agent_type = inner.agent_type;
let enter = enter_seq(agent_type);
let injected = format_injected_message(&from, &body);
drop(inner);
if let Err(e) = pty_write_paced(
&pty_writer,
agent_type,
injected.as_bytes(),
enter.as_bytes(),
) {
if let Some(id) = delivery_id {
cmd_channel.send(&Event::DeliveryFailed {
id,
reason: format!("PTY write failed: {e}"),
})?;
}
cmd_channel.send(&Event::Error {
command: "SendMessage".into(),
reason: format!("PTY write failed: {e}"),
})?;
continue;
}
let mut inner = inner_cmd.lock().unwrap();
let old = inner.state;
inner.state = ShimState::Working;
inner.state_changed_at = Instant::now();
inner.active_session_started_at = Instant::now();
let summary = inner.last_n_lines(3);
drop(inner);
if let Some(id) = delivery_id {
cmd_channel.send(&Event::MessageDelivered { id })?;
}
cmd_channel.send(&Event::StateChanged {
from: old,
to: ShimState::Working,
summary,
})?;
}
ShimState::Working => {
if inner.message_queue.len() >= MAX_QUEUE_DEPTH {
let dropped = inner.message_queue.pop_front();
let dropped_id = dropped.as_ref().and_then(|m| m.message_id.clone());
inner.message_queue.push_back(QueuedMessage {
from,
body,
message_id,
});
let depth = inner.message_queue.len();
drop(inner);
cmd_channel.send(&Event::Error {
command: "SendMessage".into(),
reason: format!(
"message queue full ({MAX_QUEUE_DEPTH}), dropped oldest message{}",
dropped_id
.map(|id| format!(" (id: {id})"))
.unwrap_or_default(),
),
})?;
cmd_channel.send(&Event::Warning {
message: format!(
"message queued while agent working (depth: {depth})"
),
idle_secs: None,
})?;
} else {
inner.message_queue.push_back(QueuedMessage {
from,
body,
message_id,
});
let depth = inner.message_queue.len();
drop(inner);
cmd_channel.send(&Event::Warning {
message: format!(
"message queued while agent working (depth: {depth})"
),
idle_secs: None,
})?;
}
}
other => {
cmd_channel.send(&Event::Error {
command: "SendMessage".into(),
reason: format!("agent in {other} state, cannot accept message"),
})?;
}
}
}
Command::CaptureScreen { last_n_lines } => {
let inner = inner_cmd.lock().unwrap();
let content = match last_n_lines {
Some(n) => inner.last_n_lines(n),
None => inner.screen_contents(),
};
let (row, col) = inner.cursor_position();
drop(inner);
cmd_channel.send(&Event::ScreenCapture {
content,
cursor_row: row,
cursor_col: col,
})?;
}
Command::GetState => {
let inner = inner_cmd.lock().unwrap();
let since = inner.state_changed_at.elapsed().as_secs();
let state = inner.state;
drop(inner);
cmd_channel.send(&Event::State {
state,
since_secs: since,
})?;
}
Command::Resize { rows, cols } => {
pty_pair
.master
.resize(PtySize {
rows,
cols,
pixel_width: 0,
pixel_height: 0,
})
.ok();
let mut inner = inner_cmd.lock().unwrap();
inner.parser.set_size(rows, cols);
}
Command::Ping => {
cmd_channel.send(&Event::Pong)?;
}
Command::Shutdown {
timeout_secs,
reason,
} => {
eprintln!(
"[shim {}] shutdown requested ({}, timeout: {}s)",
args.id,
reason.label(),
timeout_secs
);
if let Err(error) = args.preserve_work_before_kill(&args.cwd) {
eprintln!(
"[shim {}] auto-save before shutdown failed: {}",
args.id, error
);
}
{
let mut writer = pty_writer.lock().unwrap();
writer.write_all(b"\x03").ok(); writer.flush().ok();
}
let deadline = Instant::now() + Duration::from_secs(timeout_secs as u64);
loop {
if Instant::now() > deadline {
terminate_agent_group(
&mut child,
Duration::from_secs(GROUP_TERM_GRACE_SECS),
)
.ok();
break;
}
if let Ok(Some(_)) = child.try_wait() {
break;
}
thread::sleep(Duration::from_millis(PROCESS_EXIT_POLL_MS));
}
break;
}
Command::Kill => {
if let Err(error) = args.preserve_work_before_kill(&args.cwd) {
eprintln!("[shim {}] auto-save before kill failed: {}", args.id, error);
}
terminate_agent_group(&mut child, Duration::from_secs(GROUP_TERM_GRACE_SECS)).ok();
break;
}
}
}
pty_handle.join().ok();
Ok(())
}
fn drain_queue_errors(
queue: &mut VecDeque<QueuedMessage>,
terminal_state: ShimState,
) -> Vec<Event> {
common::drain_queue_errors(queue, terminal_state)
}
fn build_transition_events(
from: ShimState,
to: ShimState,
summary: &str,
pre_injection_content: &str,
current_content: &str,
last_working_screen: &str,
message_id: Option<String>,
) -> Vec<Event> {
let summary = sanitize_summary(summary);
let mut events = vec![Event::StateChanged {
from,
to,
summary: summary.clone(),
}];
if from == ShimState::Working && to == ShimState::Idle && !pre_injection_content.is_empty() {
let response =
completion_response(pre_injection_content, current_content, last_working_screen);
events.push(Event::Completion {
message_id,
response,
last_lines: summary.clone(),
});
}
if to == ShimState::ContextExhausted {
events.push(Event::ContextExhausted {
message: "Agent reported context exhaustion".to_string(),
last_lines: summary,
});
}
events
}
fn completion_response(
pre_injection_content: &str,
current_content: &str,
last_working_screen: &str,
) -> String {
let mut response = extract_response(pre_injection_content, current_content);
if response.is_empty() && !last_working_screen.is_empty() {
response = extract_response(pre_injection_content, last_working_screen);
}
response
}
fn sanitize_summary(summary: &str) -> String {
let cleaned: Vec<String> = summary
.lines()
.filter_map(|line| {
let trimmed = line.trim();
if trimmed.is_empty() || is_tui_chrome(line) || is_prompt_line(trimmed) {
return None;
}
Some(strip_claude_bullets(trimmed))
})
.collect();
if cleaned.is_empty() {
String::new()
} else {
cleaned.join("\n")
}
}
fn extract_response(pre: &str, current: &str) -> String {
let pre_lines: Vec<&str> = pre.lines().collect();
let cur_lines: Vec<&str> = current.lines().collect();
let overlap = pre_lines.len().min(cur_lines.len());
let mut diverge_at = 0;
for i in 0..overlap {
if pre_lines[i] != cur_lines[i] {
break;
}
diverge_at = i + 1;
}
let response_lines = &cur_lines[diverge_at..];
if response_lines.is_empty() {
return String::new();
}
let filtered: Vec<&str> = response_lines
.iter()
.filter(|line| !is_tui_chrome(line))
.copied()
.collect();
if filtered.is_empty() {
return String::new();
}
let mut end = filtered.len();
while end > 0 && filtered[end - 1].trim().is_empty() {
end -= 1;
}
while end > 0 && is_prompt_line(filtered[end - 1].trim()) {
end -= 1;
}
while end > 0 && filtered[end - 1].trim().is_empty() {
end -= 1;
}
let mut start = 0;
while start < end {
let trimmed = filtered[start].trim();
if trimmed.is_empty() {
start += 1;
} else if trimmed.starts_with('\u{276F}')
&& !trimmed['\u{276F}'.len_utf8()..].trim().is_empty()
{
start += 1;
} else {
break;
}
}
let cleaned: Vec<String> = filtered[start..end]
.iter()
.map(|line| strip_claude_bullets(line))
.collect();
cleaned.join("\n")
}
fn strip_claude_bullets(line: &str) -> String {
let trimmed = line.trim_start();
if trimmed.starts_with('\u{23FA}') {
let after = &trimmed['\u{23FA}'.len_utf8()..];
let leading = line.len() - line.trim_start().len();
format!("{}{}", &" ".repeat(leading), after.trim_start())
} else {
line.to_string()
}
}
fn is_tui_chrome(line: &str) -> bool {
let trimmed = line.trim();
if trimmed.is_empty() {
return false; }
if trimmed.chars().all(|c| {
matches!(
c,
'─' | '━'
| '═'
| '╌'
| '╍'
| '┄'
| '┅'
| '╶'
| '╴'
| '╸'
| '╺'
| '│'
| '┃'
| '╎'
| '╏'
| '┊'
| '┋'
)
}) {
return true;
}
if trimmed.contains("\u{23F5}\u{23F5}") || trimmed.contains("bypass permissions") {
return true;
}
if trimmed.contains("shift+tab") && trimmed.len() < 80 {
return true;
}
if trimmed.starts_with('$') && trimmed.contains("token") {
return true;
}
let braille_count = trimmed
.chars()
.filter(|c| ('\u{2800}'..='\u{28FF}').contains(c))
.count();
if braille_count > 5 {
return true;
}
let lower = trimmed.to_lowercase();
if lower.contains("welcome to the new kiro") || lower.contains("/feedback command") {
return true;
}
if lower.starts_with("kiro") && lower.contains('\u{25D4}') {
return true;
}
if trimmed.starts_with('╭') || trimmed.starts_with('╰') || trimmed.starts_with('│') {
return true;
}
if lower.starts_with("tip:") || (trimmed.starts_with('⚠') && lower.contains("limit")) {
return true;
}
if lower.contains("ask a question") || lower.contains("describe a task") {
return true;
}
false
}
fn is_prompt_line(line: &str) -> bool {
line == "\u{276F}"
|| line.starts_with("\u{276F} ")
|| line == "\u{203A}"
|| line.starts_with("\u{203A} ")
|| line.ends_with("$ ")
|| line.ends_with('$')
|| line.ends_with("% ")
|| line.ends_with('%')
|| line == ">"
|| line.starts_with("Kiro>")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn extract_response_basic() {
let pre = "line1\nline2\n$ ";
let cur = "line1\nline2\nhello world\n$ ";
assert_eq!(extract_response(pre, cur), "hello world");
}
#[test]
fn extract_response_multiline() {
let pre = "$ ";
let cur = "$ echo hi\nhi\n$ ";
let resp = extract_response(pre, cur);
assert!(resp.contains("echo hi"));
assert!(resp.contains("hi"));
}
#[test]
fn extract_response_empty() {
let pre = "$ ";
let cur = "$ ";
assert_eq!(extract_response(pre, cur), "");
}
#[test]
fn content_hash_deterministic() {
assert_eq!(content_hash("hello"), content_hash("hello"));
assert_ne!(content_hash("hello"), content_hash("world"));
}
#[test]
fn shell_single_quote_escapes_embedded_quote() {
assert_eq!(shell_single_quote("fix user's bug"), "fix user'\\''s bug");
}
#[test]
fn supervised_command_contains_watchdog_and_exec() {
let command = build_supervised_agent_command("kiro-cli chat 'hello'", 4242);
assert!(command.contains("shim_pid=4242"));
assert!(command.contains("agent_root_pid=$$"));
assert!(command.contains("agent_pgid=$$"));
assert!(command.contains("setsid sh -c"));
assert!(command.contains("shim_pid=\"$1\""));
assert!(command.contains("agent_pgid=\"$2\""));
assert!(command.contains("agent_root_pid=\"$3\""));
assert!(command.contains("collect_descendants()"));
assert!(command.contains("pgrep -P \"$parent_pid\""));
assert!(command.contains("descendant_pids=$(collect_descendants \"$agent_root_pid\")"));
assert!(command.contains("kill -TERM -- -\"$agent_pgid\""));
assert!(command.contains("kill -TERM \"$descendant_pid\""));
assert!(command.contains("kill -KILL -- -\"$agent_pgid\""));
assert!(command.contains("kill -KILL \"$descendant_pid\""));
assert!(command.contains("' _ \"$shim_pid\" \"$agent_pgid\" \"$agent_root_pid\""));
assert!(command.contains("exec bash -lc 'kiro-cli chat '\\''hello'\\'''"));
}
#[test]
fn is_prompt_line_shell_dollar() {
assert!(is_prompt_line("user@host:~$ "));
assert!(is_prompt_line("$"));
}
#[test]
fn is_prompt_line_claude() {
assert!(is_prompt_line("\u{276F}"));
assert!(is_prompt_line("\u{276F} "));
}
#[test]
fn is_prompt_line_codex() {
assert!(is_prompt_line("\u{203A}"));
assert!(is_prompt_line("\u{203A} "));
}
#[test]
fn is_prompt_line_kiro() {
assert!(is_prompt_line("Kiro>"));
assert!(is_prompt_line(">"));
}
#[test]
fn is_prompt_line_not_prompt() {
assert!(!is_prompt_line("hello world"));
assert!(!is_prompt_line("some output here"));
}
#[test]
fn build_transition_events_working_to_idle() {
let events = build_transition_events(
ShimState::Working,
ShimState::Idle,
"summary",
"pre\n$ ",
"pre\nhello\n$ ",
"",
Some("msg-1".into()),
);
assert_eq!(events.len(), 2);
assert!(matches!(&events[0], Event::StateChanged { .. }));
assert!(matches!(&events[1], Event::Completion { .. }));
}
#[test]
fn completion_response_uses_last_working_screen_fallback() {
let response = completion_response("pre\n$ ", "pre\n$ ", "pre\nfailed tests\n$ ");
assert_eq!(response, "failed tests");
}
#[test]
fn build_transition_events_to_context_exhausted() {
let events = build_transition_events(
ShimState::Working,
ShimState::ContextExhausted,
"summary",
"",
"",
"",
None,
);
assert_eq!(events.len(), 2);
assert!(matches!(&events[1], Event::ContextExhausted { .. }));
}
#[test]
fn build_transition_events_starting_to_idle() {
let events = build_transition_events(
ShimState::Starting,
ShimState::Idle,
"summary",
"",
"",
"",
None,
);
assert_eq!(events.len(), 1);
assert!(matches!(&events[0], Event::StateChanged { .. }));
}
fn make_queued_msg(id: &str, body: &str) -> QueuedMessage {
QueuedMessage {
from: "user".into(),
body: body.into(),
message_id: Some(id.into()),
}
}
#[test]
fn queue_enqueue_basic() {
let mut queue: VecDeque<QueuedMessage> = VecDeque::new();
queue.push_back(make_queued_msg("m1", "hello"));
queue.push_back(make_queued_msg("m2", "world"));
assert_eq!(queue.len(), 2);
}
#[test]
fn queue_fifo_order() {
let mut queue: VecDeque<QueuedMessage> = VecDeque::new();
queue.push_back(make_queued_msg("m1", "first"));
queue.push_back(make_queued_msg("m2", "second"));
queue.push_back(make_queued_msg("m3", "third"));
let msg = queue.pop_front().unwrap();
assert_eq!(msg.message_id.as_deref(), Some("m1"));
assert_eq!(msg.body, "first");
let msg = queue.pop_front().unwrap();
assert_eq!(msg.message_id.as_deref(), Some("m2"));
assert_eq!(msg.body, "second");
let msg = queue.pop_front().unwrap();
assert_eq!(msg.message_id.as_deref(), Some("m3"));
assert_eq!(msg.body, "third");
assert!(queue.is_empty());
}
#[test]
fn queue_overflow_drops_oldest() {
let mut queue: VecDeque<QueuedMessage> = VecDeque::new();
for i in 0..MAX_QUEUE_DEPTH {
queue.push_back(make_queued_msg(&format!("m{i}"), &format!("msg {i}")));
}
assert_eq!(queue.len(), MAX_QUEUE_DEPTH);
assert!(queue.len() >= MAX_QUEUE_DEPTH);
let dropped = queue.pop_front().unwrap();
assert_eq!(dropped.message_id.as_deref(), Some("m0")); queue.push_back(make_queued_msg("m_new", "new message"));
assert_eq!(queue.len(), MAX_QUEUE_DEPTH);
let first = queue.pop_front().unwrap();
assert_eq!(first.message_id.as_deref(), Some("m1"));
}
#[test]
fn drain_queue_errors_empty() {
let mut queue: VecDeque<QueuedMessage> = VecDeque::new();
let events = drain_queue_errors(&mut queue, ShimState::Dead);
assert!(events.is_empty());
}
#[test]
fn drain_queue_errors_with_messages() {
let mut queue: VecDeque<QueuedMessage> = VecDeque::new();
queue.push_back(make_queued_msg("m1", "hello"));
queue.push_back(make_queued_msg("m2", "world"));
queue.push_back(QueuedMessage {
from: "user".into(),
body: "no id".into(),
message_id: None,
});
let events = drain_queue_errors(&mut queue, ShimState::Dead);
assert_eq!(events.len(), 3);
assert!(queue.is_empty());
for event in &events {
assert!(matches!(event, Event::Error { .. }));
}
if let Event::Error { reason, .. } = &events[0] {
assert!(reason.contains("dead"));
assert!(reason.contains("m1"));
}
if let Event::Error { reason, .. } = &events[2] {
assert!(!reason.contains("(id:"));
}
}
#[test]
fn drain_queue_errors_context_exhausted() {
let mut queue: VecDeque<QueuedMessage> = VecDeque::new();
queue.push_back(make_queued_msg("m1", "hello"));
let events = drain_queue_errors(&mut queue, ShimState::ContextExhausted);
assert_eq!(events.len(), 1);
if let Event::Error { reason, .. } = &events[0] {
assert!(reason.contains("context_exhausted"));
}
}
#[test]
fn queued_message_preserves_fields() {
let msg = QueuedMessage {
from: "manager".into(),
body: "do this task".into(),
message_id: Some("msg-42".into()),
};
assert_eq!(msg.from, "manager");
assert_eq!(msg.body, "do this task");
assert_eq!(msg.message_id.as_deref(), Some("msg-42"));
}
#[test]
fn queued_message_none_id() {
let msg = QueuedMessage {
from: "user".into(),
body: "anonymous".into(),
message_id: None,
};
assert!(msg.message_id.is_none());
}
#[test]
fn max_queue_depth_is_16() {
assert_eq!(MAX_QUEUE_DEPTH, 16);
}
#[test]
fn format_injected_message_includes_sender_and_reply_target() {
let formatted = format_injected_message("human", "what is 2+2?");
assert!(formatted.contains("--- Message from human ---"));
assert!(formatted.contains("Reply-To: human"));
assert!(formatted.contains("batty send human"));
assert!(formatted.ends_with("what is 2+2?"));
}
#[test]
fn format_injected_message_uses_sender_as_reply_target() {
let formatted = format_injected_message("manager", "status?");
assert!(formatted.contains("Reply-To: manager"));
assert!(formatted.contains("batty send manager"));
}
#[test]
fn sanitize_summary_strips_tui_chrome_and_prompt_lines() {
let summary = "────────────────────\n❯ \n ⏵⏵ bypass permissions on\nThe answer is 4\n";
assert_eq!(sanitize_summary(summary), "The answer is 4");
}
#[test]
fn sanitize_summary_keeps_multiline_meaningful_content() {
let summary = " Root cause: stale resume id\n\n Fix: retry with fresh start\n";
assert_eq!(
sanitize_summary(summary),
"Root cause: stale resume id\nFix: retry with fresh start"
);
}
#[test]
fn is_tui_chrome_horizontal_rule() {
assert!(is_tui_chrome("────────────────────────────────────"));
assert!(is_tui_chrome(" ───────── "));
assert!(is_tui_chrome("━━━━━━━━━━━━━━━━━━━━"));
}
#[test]
fn is_tui_chrome_status_bar() {
assert!(is_tui_chrome(
" \u{23F5}\u{23F5} bypass permissions on (shift+tab to toggle)"
));
assert!(is_tui_chrome(" bypass permissions on"));
assert!(is_tui_chrome(" shift+tab"));
}
#[test]
fn is_tui_chrome_cost_line() {
assert!(is_tui_chrome("$0.01 · 2.3k tokens"));
}
#[test]
fn is_tui_chrome_not_content() {
assert!(!is_tui_chrome("Hello, world!"));
assert!(!is_tui_chrome("The answer is 4"));
assert!(!is_tui_chrome("")); assert!(!is_tui_chrome(" some output "));
}
#[test]
fn extract_response_strips_chrome() {
let pre = "idle screen\n\u{276F} ";
let cur = "\u{276F} Hello\n\nThe answer is 42\n\n\
────────────────────\n\
\u{23F5}\u{23F5} bypass permissions on\n\
\u{276F} ";
let resp = extract_response(pre, cur);
assert!(resp.contains("42"), "should contain the answer: {resp}");
assert!(
!resp.contains("────"),
"should strip horizontal rule: {resp}"
);
assert!(!resp.contains("bypass"), "should strip status bar: {resp}");
}
#[test]
fn extract_response_strips_echoed_input() {
let pre = "\u{276F} ";
let cur = "\u{276F} What is 2+2?\n\n4\n\n\u{276F} ";
let resp = extract_response(pre, cur);
assert!(resp.contains('4'), "should contain answer: {resp}");
assert!(
!resp.contains("What is 2+2"),
"should strip echoed input: {resp}"
);
}
#[test]
fn extract_response_tui_full_rewrite() {
let pre = "Welcome to Claude\n\n\u{276F} ";
let cur = "\u{276F} Hello\n\nHello! How can I help?\n\n\
────────────────────\n\
\u{276F} ";
let resp = extract_response(pre, cur);
assert!(
resp.contains("Hello! How can I help?"),
"should extract response from TUI rewrite: {resp}"
);
}
#[test]
fn strip_claude_bullets_removes_marker() {
assert_eq!(strip_claude_bullets("\u{23FA} 4"), "4");
assert_eq!(
strip_claude_bullets(" \u{23FA} hello world"),
" hello world"
);
assert_eq!(strip_claude_bullets("no bullet here"), "no bullet here");
assert_eq!(strip_claude_bullets(""), "");
}
#[test]
fn extract_response_strips_claude_bullets() {
let pre = "\u{276F} ";
let cur = "\u{276F} question\n\n\u{23FA} 42\n\n\u{276F} ";
let resp = extract_response(pre, cur);
assert!(resp.contains("42"), "should contain answer: {resp}");
assert!(
!resp.contains('\u{23FA}'),
"should strip bullet marker: {resp}"
);
}
#[test]
fn preserve_handoff_writes_diff_and_commit_summary() {
let repo = tempfile::tempdir().unwrap();
init_test_git_repo(repo.path());
let events_dir = repo.path().join(".batty").join("team_config");
std::fs::create_dir_all(&events_dir).unwrap();
std::fs::write(repo.path().join("tracked.txt"), "one\n").unwrap();
run_test_git(repo.path(), &["add", "tracked.txt"]);
run_test_git(repo.path(), &["commit", "-m", "initial commit"]);
std::fs::write(repo.path().join("tracked.txt"), "one\ntwo\n").unwrap();
let task_id = "42";
let event_lines = [
serde_json::to_string(&crate::team::events::TeamEvent::task_assigned(
"eng-1", task_id,
))
.unwrap(),
serde_json::to_string(&crate::team::events::TeamEvent::agent_restarted(
"eng-1",
task_id,
"context_exhausted",
1,
))
.unwrap(),
];
std::fs::write(events_dir.join("events.jsonl"), event_lines.join("\n")).unwrap();
let recent_output = "\
running cargo test --lib\n\
test result: ok\n\
editing src/lib.rs\n";
let task = crate::task::Task {
id: 42,
title: "resume widget".to_string(),
status: "in-progress".to_string(),
priority: "high".to_string(),
assignee: None,
claimed_by: Some("eng-1".to_string()),
claimed_at: None,
claim_ttl_secs: None,
claim_expires_at: None,
last_progress_at: None,
claim_warning_sent_at: None,
claim_extensions: None,
last_output_bytes: None,
blocked: None,
tags: Vec::new(),
depends_on: Vec::new(),
review_owner: None,
blocked_on: None,
worktree_path: None,
branch: None,
commit: None,
artifacts: Vec::new(),
next_action: Some("Run the Rust tests and finish the restart handoff.".to_string()),
scheduled_for: None,
cron_schedule: None,
cron_last_run: None,
completed: None,
description: "Continue widget implementation.".to_string(),
batty_config: None,
source_path: repo.path().join("task-42.md"),
};
preserve_handoff(repo.path(), &task, Some(recent_output)).unwrap();
let handoff = std::fs::read_to_string(repo.path().join(HANDOFF_FILE_NAME)).unwrap();
assert!(handoff.contains("# Carry-Forward Summary"));
assert!(handoff.contains("## Task Spec"));
assert!(handoff.contains("Task #42: resume widget"));
assert!(handoff.contains("## Work Completed So Far"));
assert!(handoff.contains("### Branch"));
assert!(handoff.contains("master") || handoff.contains("main"));
assert!(handoff.contains("### Last Commit"));
assert!(handoff.contains("### Changed Files"));
assert!(handoff.contains("tracked.txt"));
assert!(handoff.contains("### Tests Run"));
assert!(handoff.contains("cargo test --lib"));
assert!(handoff.contains("### Progress Summary"));
assert!(handoff.contains("task_assigned"));
assert!(handoff.contains("agent_restarted: context_exhausted"));
assert!(handoff.contains("### Recent Activity"));
assert!(handoff.contains("editing src/lib.rs"));
assert!(handoff.contains("### Recent Commits"));
assert!(handoff.contains("initial commit"));
assert!(handoff.contains("## What Remains"));
assert!(handoff.contains("Run the Rust tests and finish the restart handoff."));
}
#[test]
fn preserve_handoff_uses_none_when_repo_has_no_changes_or_commits() {
let repo = tempfile::tempdir().unwrap();
init_test_git_repo(repo.path());
let task = crate::task::Task {
id: 7,
title: "empty repo".to_string(),
status: "in-progress".to_string(),
priority: "low".to_string(),
assignee: None,
claimed_by: Some("eng-1".to_string()),
claimed_at: None,
claim_ttl_secs: None,
claim_expires_at: None,
last_progress_at: None,
claim_warning_sent_at: None,
claim_extensions: None,
last_output_bytes: None,
blocked: None,
tags: Vec::new(),
depends_on: Vec::new(),
review_owner: None,
blocked_on: None,
worktree_path: None,
branch: None,
commit: None,
artifacts: Vec::new(),
next_action: None,
scheduled_for: None,
cron_schedule: None,
cron_last_run: None,
completed: None,
description: "No changes yet.".to_string(),
batty_config: None,
source_path: repo.path().join("task-7.md"),
};
preserve_handoff(repo.path(), &task, None).unwrap();
let handoff = std::fs::read_to_string(repo.path().join(HANDOFF_FILE_NAME)).unwrap();
assert!(handoff.contains("### Changed Files\n(none)"));
assert!(handoff.contains("### Tests Run\n(none)"));
assert!(handoff.contains("### Recent Activity\n(none)"));
assert!(handoff.contains("### Recent Commits\n(none)"));
assert!(handoff.contains("## What Remains"));
}
#[test]
fn extract_test_commands_deduplicates_known_test_invocations() {
let output = "\
\u{1b}[31mcargo test --lib\u{1b}[0m\n\
pytest tests/test_api.py\n\
cargo test --lib\n\
plain output\n";
let tests = extract_test_commands(output);
assert_eq!(
tests,
vec![
"cargo test --lib".to_string(),
"pytest tests/test_api.py".to_string()
]
);
}
#[test]
fn preserve_work_before_kill_respects_config_toggle() {
let tmp = tempfile::tempdir().unwrap();
let preserved =
preserve_work_before_kill_with(tmp.path(), Duration::from_millis(10), false, |_path| {
panic!("commit should not run when disabled")
})
.unwrap();
assert!(!preserved);
}
#[test]
fn preserve_work_before_kill_times_out() {
let tmp = tempfile::tempdir().unwrap();
let (_tx, rx) = std::sync::mpsc::channel::<()>();
let preserved = preserve_work_before_kill_with(
tmp.path(),
Duration::from_millis(10),
true,
move |_path| {
let _ = rx.recv_timeout(Duration::from_millis(100));
Ok(true)
},
)
.unwrap();
assert!(!preserved);
}
fn init_test_git_repo(path: &Path) {
run_test_git(path, &["init"]);
run_test_git(path, &["config", "user.name", "Batty Tests"]);
run_test_git(path, &["config", "user.email", "batty-tests@example.com"]);
}
fn run_test_git(path: &Path, args: &[&str]) {
use std::process::Command;
let output = Command::new("git")
.args(args)
.current_dir(path)
.output()
.unwrap();
assert!(
output.status.success(),
"git {} failed: {}",
args.join(" "),
String::from_utf8_lossy(&output.stderr)
);
}
}